diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 1dc082947..2f5b05e6e 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -26,10 +26,29 @@ topics: - arena - physics step - sim update + - runtime control + - kinematic joint trajectory + - register_kinematic_joint_trajectory + - kinematic nodal trajectory + - register_kinematic_nodal_trajectory + - KinematicJointTrajectoryControl + - contact material schedule + - register_contact_material_schedule + - particle contact material schedule + - register_particle_contact_material_schedule - scene object - asset registry - manual update - GPU physics + - deformable + - particle set + - Newton soft body + - Newton cloth + - array-backed mesh + - stable node indices + - tetrahedral particle order + - independent visual mesh + - visual binding mode - simulation lifecycle - collision_policy - collision isolation @@ -53,10 +72,12 @@ topics: source_of_truth: - embodichain/lab/sim/__init__.py - embodichain/lab/sim/sim_manager.py + - embodichain/lab/sim/_runtime_controls.py - embodichain/lab/sim/cfg/ - embodichain/lab/sim/_legacy_cfg.py - embodichain/lab/sim/common.py - embodichain/lab/sim/material.py + - embodichain/lab/sim/shapes.py - embodichain/lab/sim/profiler.py - embodichain/lab/sim/spawn/descriptors.py - embodichain/lab/sim/spawn/scene.py @@ -361,6 +382,9 @@ topics: - deformable - soft body - cloth + - Newton deformable + - particle set + - render topology - quaternion - xyzw paths: @@ -594,6 +618,7 @@ topics: title: Differentiable Environment (APG) aliases: - differentiable env + - DifferentiableEnv - apg - analytic policy gradient - differentiable rl @@ -608,7 +633,7 @@ topics: - warp tape - requires_grad - semi_implicit - - DifferentiableEmbodiedEnv + - DifferentiableEnv - NewtonStepFunc - quaternion - xyzw @@ -616,6 +641,7 @@ topics: - topics/differentiable-env/differentiable-env.md source_of_truth: - embodichain/lab/gym/envs/differentiable_env.py + - embodichain/lab/sim/cfg/simulation.py - embodichain/lab/sim/diff/ - embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py related_topics: diff --git a/agent_context/topics/differentiable-env/differentiable-env.md b/agent_context/topics/differentiable-env/differentiable-env.md index dc58718b3..4a14a4161 100644 --- a/agent_context/topics/differentiable-env/differentiable-env.md +++ b/agent_context/topics/differentiable-env/differentiable-env.md @@ -1,114 +1,95 @@ # differentiable-env -> Topic: Differentiable environment for analytic policy gradient (APG) — -> `DifferentiableEmbodiedEnv` + the `embodichain.lab.sim.diff` Warp-tape -> ↔ PyTorch-autograd bridge. +> Topic: Newton-backed kinematic environments for analytic policy gradient +> (APG) through the Warp-tape ↔ PyTorch-autograd bridge. + +## Public entry point + +Use +`embodichain.lab.gym.envs.differentiable_env.DifferentiableEnv`. +It inherits `EmbodiedEnv`, preserves its scene lifecycle, and replaces the +normal physics step with a task-defined kinematics callback recorded on a +Warp tape. + +The resolution path is: + + DifferentiableEnv.step(action) + → NewtonStepFunc.apply(action, sim_state) + → _apply_action_kernel(action_wp, tape) + → _make_kinematic_step_fn()() + → _read_outputs(final_state) + → Warp tape backward → action.grad + +## Invariants + +- The configured physics backend must be `NewtonPhysicsCfg`. +- `NewtonPhysicsCfg.requires_grad` must be `True`. +- Default physics and other backends fail during `DifferentiableEnv` + construction. +- `DifferentiableEnv` always supplies a named kinematics callback; its bridge + contract has no dynamics mode, solver substeps, or control buffer. +- The environment never invokes the configured Newton solver or collision + pipeline. +- Newton gradient configuration still selects the semi-implicit solver, but + `DifferentiableEnv` never advances it. +- Gradient mode disables Newton CUDA graph capture. -## Overview +## Subclass contract -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()`. +Task authors implement three hooks: -## Required configuration +- `_apply_action_kernel(action_wp, tape)` launches Warp work that maps the + PyTorch action bridge array into task-owned kinematic state. +- `_make_kinematic_step_fn()` returns a zero-argument callback such as + `newton.eval_fk(...)`. The callback returns the state consumed by the output + hook. +- `_read_outputs(final_state)` returns `obs`, `reward`, `terminated`, and + `truncated`, plus `_order` and `_grad_track` metadata used by + `NewtonStepFunc`. -- `NewtonPhysicsCfg(requires_grad=True, solver_cfg={"solver_type": "semi_implicit"})` -- `use_cuda_graph=False` (forced by dexsim when grad mode is on) +There is no public `_apply_dynamics_action_kernel` or +`differentiable_step_mode` extension point. Differentiable dynamics are a +future feature, not a `DifferentiableEnv` capability. -The default backend and any other Newton solver are rejected at -construction time by `DifferentiableEmbodiedEnv._validate_diff_cfg`. +## Autograd and reset rules -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)`. +Action, kinematics, and gradient-producing output kernels must execute while +the Warp tape is open. Each tracked output names its backing Warp array in +`_grad_track`; an output mapped to `None` does not seed Warp backward. -## Subclass contract +A grad-tracked terminal step returns the terminal observation and exposes +`requires_reset_after_backward` plus `deferred_reset_ids` in `info`. Reset +those rows only after backward. A no-grad terminal step resets them +synchronously. + +## Franka reference task + +`embodichain_tasks.special.franka_reach_apg.FrankaReachApgEnv` is the canonical +example. Its path is: + + action → new_joint_q → newton.eval_fk → body_q → reward kernel → action.grad -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. +The task snapshots live joint positions before opening the tape and writes the +detached next joint state back after the bridge returns. It does not exercise +Newton dynamics. + +## Dynamics boundary + +The public differentiable package exposes no solver stepper, trajectory, or +gradient-rollout API. Add those capabilities as a separate future design when +Newton dynamics are ready for end-to-end validation. ## 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. +- `embodichain/lab/gym/envs/differentiable_env.py` +- `embodichain/lab/sim/cfg/simulation.py` +- `embodichain/lab/sim/diff/` +- `embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py` + +## Focused validation + +- `tests/gym/envs/test_differentiable_embodied_env.py` +- `tests/sim/test_sim_manager_cfg.py` ## Related topics diff --git a/agent_context/topics/sim-visualization/sim-visualization.md b/agent_context/topics/sim-visualization/sim-visualization.md index 58ae1dfa8..3c4c74a8d 100644 --- a/agent_context/topics/sim-visualization/sim-visualization.md +++ b/agent_context/topics/sim-visualization/sim-visualization.md @@ -206,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 | -| 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 | +| Volume `DeformableObject` (`SoftObject`) | Live Newton render-surface vertices and triangles | +| Surface `DeformableObject` (`ClothObject`) | Live Newton render-surface vertices and triangles | | `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 | @@ -235,21 +235,28 @@ slow rendering or clients cannot accumulate an image backlog. ## Deformables -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. - `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 - with `cKDTree`. Construction raises `RuntimeError` if the mapping distance - exceeds the scale-relative tolerance. +Volume and surface deformables require the Newton backend, CUDA, and a +particle-capable solver. Their live vertices are sampled at `soft_body_fps`, +independently from `scene_fps`. `SceneExporter` enumerates the manager's single +deformable registry and reads render topology through +`get_surface_vertices()` and `get_surface_triangles()`; it does not branch on +legacy buffer APIs. The facade returns world-frame render vertices, and the +exporter subtracts the arena offset before publishing them below the arena +node. The `deformable_type` discriminator only selects the existing +soft/cloth browser node kind, path, and color. + +- Both soft bodies and cloth publish the live render surface exposed by their + DexSim 0.5 typed Newton particle-set handles. No convex-hull reconstruction + or nearest-neighbor welding is performed in the visualization path. +- A volume deformable separately exposes its tetrahedral collision surface + through `get_collision_surface_triangles()` for consumers that need physical + rather than render topology. `SceneExporter` intentionally uses render + topology. +- Spawn binding validates that every replicated instance has the same render + vertex and triangle counts. A file-backed soft body whose DexSim clones have + a render topology different from the source fails during scene preparation; + use one environment or a compatible mesh until DexSim replication preserves + the source topology. - Viser does not update mesh vertices in place. `ViserBackend` removes and recreates a deformable mesh handle only when a dynamic vertex sample arrives. Pose-only frames reuse the current handle. @@ -318,8 +325,7 @@ payload bytes plus capture/upload time. | Startup timeout or address-in-use error | The Viser worker did not become ready or the configured port is occupied. Select another port and inspect `visualization_health.worker_error`. | | Asset added after startup is missing | Step once, call `refresh_visualization()`, or mark topology dirty if the change bypassed manager APIs. | | Browser stops updating after an exporter/backend exception | `capture_visualization_safely()` latches the first error to protect simulation. Inspect health/logs, then stop and restart after fixing the cause. | -| Soft body looks inflated or loses cavities | The surface is a collision-vertex convex hull, not the render topology. | -| Cloth construction raises a mapping error | Render vertices do not match the welded physical rest vertices within tolerance. | +| Replicated soft body fails with a render-vertex-count mismatch | DexSim produced clone render topology different from the source. Use one environment or a compatible mesh until DexSim replication preserves the topology. | | Camera frustum exists but preview is blank | Color capture is disabled, no image has been captured yet, or the selected camera/environment is hidden. | | Stereo/contact sensor is absent | Current camera export accepts only `sensor_type == "Camera"`; non-mesh sensors are not exported. | | Browser lags or upload cost is high | Reduce scene/image/deformable FPS, select fewer environments, or lower point-cloud limits. | diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 82a9182da..e83434311 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -25,7 +25,16 @@ 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. +place. Newton kinematic trajectory controls are registered through +`SimulationManager.register_kinematic_joint_trajectory()` and +`SimulationManager.register_kinematic_nodal_trajectory()`; the manager expands +the arena batch to concrete Spawn paths without exposing its private +`SpawnScene` or `SceneBuilder`. Nodal trajectories move configured inactive +deformable particles relative to their finalized initial positions at every +Newton substep. Time-varying rigid/articulation and scene-wide particle contact +properties use `register_contact_material_schedule()` and +`register_particle_contact_material_schedule()` at the same pre-`prepare()` +boundary. The registries cover: @@ -53,6 +62,7 @@ EnvCfg.sim_cfg → EmbodiedEnv declares robot, objects, lights, and physical sensors → Default may materialize native handles eagerly → Newton keeps physical descriptors deferred + → optionally register Newton trajectory or contact-material controls on the manager → SimulationManager.prepare() → for Newton, resolve source metadata and configure exact-name overlays → finalize/rebuild pending Spawn descriptors once @@ -117,21 +127,25 @@ 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`, +`DeformableObject`/`DeformableObjectData` contract and the Newton particle-set +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`. +`SoftBodyDesc` and `ClothDesc` particle-set descriptors; 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. +Only the Newton backend is registered; the Default backend intentionally +reports both deformable capabilities as unsupported. Declaration requires CUDA +and a particle-capable Newton solver (`xpbd`, `semi_implicit`, `vbd`, or +`mjvbd`), rejects gradient mode and post-finalization additions, and compiles +directly to DexSim 0.5 `SoftBodyDesc` or `ClothDesc`. Runtime state is fetched +and applied through `Scene.create_particle_set_batch()`; direct DexSim +`SoftBody`/`ClothBody` buffers and the old utility loaders are not supported. +Volume collision topology comes from the typed particle-set handle, while +render topology and vertices remain a separate visualization surface. `BaseEnv._setup_scene()` temporarily constructs the manager headlessly so the scene can be assembled before a native window is opened. It sets @@ -179,10 +193,39 @@ reconfigured or rebound. `init_gpu_physics()` and 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. +reading link/joint metadata, object state, or advancing physics. A Newton +caller that needs a substep-interpolated kinematic articulation must call +`register_kinematic_joint_trajectory(uid, joint_positions, ...)` after +declaring that robot/articulation and before `prepare()`. Its public trajectory +layout is `[num_envs, frames, dof]` in the articulation's public qpos order; +the optional root-pose layout is +`[num_envs, frames, 4, 4]`. `BaseEnv` provides the readiness boundary +automatically between `_setup_scene()` and metadata-dependent setup. +`SimulationManager.update()` still calls the readiness path defensively before +advancing the requested physics steps. + +A caller that needs selected deformable nodes to follow a kinematic path must +clear the Newton `ACTIVE` bit for those nodes through +`DeformableObjectCfg.particle_flags`, then call +`register_kinematic_nodal_trajectory(uid, node_indices, position_offsets, ...)` +before `prepare()`. Offsets use the batched layout +`[num_envs, samples, selected_nodes, 3]` and are relative to the world positions +captured when the finalized runtime control initializes. With `fps`, targets +are linearly interpolated at substep times; without it, each substep consumes +one sample. Surface indices can follow an array-backed source mesh directly. +Volume indices address the generated tetrahedral simulation particles and do +not correspond to source-mesh vertices. The host-side particle writes +intentionally disable CUDA Graph replay for that scene. + +Time-varying Newton shape contact properties belong on +`register_contact_material_schedule(uid, keyframes, ...)`; valid targets are a +declared rigid object, robot, or articulation. The manager creates one control +per Arena and accepts piecewise-constant `dynamic_friction`, `stiffness`, and +`damping` tracks. Scene-wide particle-versus-rigid values use +`register_particle_contact_material_schedule(keyframes)`. That control performs +host-side writes and disables CUDA Graph replay, while shape-material and joint +trajectory controls remain graph-compatible. Register all controls before +`prepare()` and never reach into `_spawn_scene.builder` from a task or demo. ## Module Boundaries @@ -193,7 +236,7 @@ readiness path defensively before advancing the requested physics steps. | 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 | +| Common deformable contract and Newton particle-set 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` | @@ -249,9 +292,27 @@ module or the corresponding robot/sensor module. Scene composition belongs in 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. +the volume subclass, and Newton triangle/edge/spring attributes stay on the +surface subclass. `particle_radius` and `validate_mesh` are common particle-set +options. `particle_flags` accepts one Newton bitmask or a +simulation-particle-order array; nodes driven by a kinematic nodal trajectory +must have their `ACTIVE` bit cleared before the immutable solver model is +constructed. `MeshCfg` accepts either a file path or explicit vertex/triangle +arrays. For surface deformables, use the array-backed form when flags or +trajectories require stable source node indices: native file importers may +reorder vertices, while array-backed descriptors preserve the supplied order +through Newton model construction. A surface deformable may additionally set +`visual_shape` to an independently indexed render mesh and choose +`visual_binding_mode="auto"` or `"nearest_vertex"`; physics always follows +`shape`, and visual material/UV configuration belongs on `visual_shape` when it +is present. This supports seam-duplicated textured meshes without changing +particle topology. Volume deformables are voxelized into a +separate tetrahedral simulation mesh, so their particle indices are not source +mesh indices. The volume descriptor converts Young's modulus and Poisson's +ratio to Newton Lamé coefficients. Removed Default-only fields are deliberately +not translated or silently ignored. Do not add backend conditionals to one +monolithic deformable config; a future implementation belongs at the manager +dispatch boundary. New rigid-body configs use `RigidBodyPhysicsCfg`. Portable intent is organized by physical concept: @@ -465,9 +526,27 @@ legacy layer can eventually be removed as one unit. - 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. +- Create deformables only with the Newton backend on CUDA and a supported + particle solver; Default-backend soft/cloth compatibility is intentionally + absent. +- Treat deformable render meshes and physical particle topology as distinct; + state mutation uses the particle batch and never writes renderer/native + soft-body buffers directly. - 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. +- Register Newton kinematic joint trajectories through + `SimulationManager.register_kinematic_joint_trajectory()` before `prepare()`; + callers must not access `SimulationManager._spawn_scene` or its builder. +- Register deformable kinematic-node trajectories through + `SimulationManager.register_kinematic_nodal_trajectory()` before `prepare()`; + mark every selected node inactive in `particle_flags` before model + construction, and keep Spawn/runtime-control access inside the manager. +- Register rigid/articulation contact schedules through + `SimulationManager.register_contact_material_schedule()` and scene-wide + particle schedules through + `SimulationManager.register_particle_contact_material_schedule()` before + `prepare()`; account for the latter disabling CUDA Graph replay. - 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 @@ -496,6 +575,8 @@ legacy layer can eventually be removed as one unit. | Scene resource cannot be found or the wrong object is returned | UID mismatch or code bypassed the manager registry | | 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 | +| Adding a soft or cloth object fails immediately on the Default backend | Deformables are Newton-only; select a supported Newton particle solver on CUDA | +| A replicated file-backed soft body fails render upload because clone vertex counts differ | DexSim's cloned render mesh does not match the template embedding topology; use a compatible mesh/single environment while the DexSim 0.5 clone path is corrected | | 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 | diff --git a/design/newton-backend-design.md b/design/newton-backend-design.md index f809e8156..80d88631e 100644 --- a/design/newton-backend-design.md +++ b/design/newton-backend-design.md @@ -337,26 +337,22 @@ Done: `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: + (`torch.autograd.Function`) bridging task-defined Newton kinematics + recorded on a `wp.Tape` into PyTorch autograd. `DifferentiableEnv` + validates `NewtonPhysicsCfg(requires_grad=True)` and never advances the + configured semi-implicit solver. The Franka FR3 + reach APG example (`franka_reach_apg.py`) exercises the bridge end-to-end + with `newton.eval_fk`, a Warp action kernel, and a Warp reward kernel + computed inside the tape. Solver steppers, differentiable trajectories, + and gradient rollouts are not part of this kinematics-only stage; they need + a separate future dynamics design. 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. + The Franka task uses a task-defined FK callback + (``newton.eval_fk``). This is the only supported differentiable + environment route in the current stage; the configured + ``semi_implicit`` solver is not advanced. Remaining: 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 107a90955..024d29655 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst @@ -60,14 +60,15 @@ Environment Classes 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. +``DifferentiableEnv`` keeps the standard environment lifecycle while bridging +task-defined Newton kinematics into PyTorch autograd for analytic +policy-gradient tasks. Subclasses provide action, kinematics, and output +kernels; the base class owns tape-aware stepping and deferred resets without +advancing the Newton solver. .. currentmodule:: embodichain.lab.gym.envs.differentiable_env -.. autoclass:: DifferentiableEmbodiedEnv +.. autoclass:: DifferentiableEnv :members: :inherited-members: :show-inheritance: 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 b32ba8f07..ca8bcc295 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst @@ -23,6 +23,9 @@ 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. +Surface deformables may keep a stable low-resolution simulation topology in +``shape`` while binding an independently indexed ``visual_shape`` for authored +UV seams and render detail. .. rubric:: Classes 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..986062ea2 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.shapes.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.shapes.rst @@ -6,11 +6,14 @@ embodichain.lab.sim.shapes Overview -------- -Geometry configuration objects used to build the collision and visual shapes of -rigid bodies. :class:`ShapeCfg` is the common base; :class:`MeshCfg`, +Geometry configuration objects used to build collision, visual, and deformable +simulation shapes. :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. +assets are loaded and decomposed. ``MeshCfg`` accepts either a file path or +explicit vertex/triangle arrays; for surface deformables, the array-backed form +preserves node order for particle flags and kinematic trajectories. Volume +deformables generate a separate tetrahedral simulation mesh during voxelization. .. rubric:: Classes 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 611539f88..dfb114798 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 @@ -27,7 +27,23 @@ instance registry instead of passing it around explicitly. :members: :undoc-members: :show-inheritance: - :exclude-members: visualize_point_cloud + :exclude-members: register_contact_material_schedule, register_kinematic_joint_trajectory, register_kinematic_nodal_trajectory, register_particle_contact_material_schedule, visualize_point_cloud + +.. rubric:: Newton runtime controls + +Runtime controls must be registered after declaring their target assets and +before :meth:`SimulationManager.prepare`. The manager expands logical UIDs to +the concrete paths of every Arena, so callers do not need access to the private +Spawn scene. The particle-material schedule is host-side and disables CUDA +Graph replay; the other controls are graph-compatible. + +.. automethod:: SimulationManager.register_kinematic_joint_trajectory + +.. automethod:: SimulationManager.register_kinematic_nodal_trajectory + +.. automethod:: SimulationManager.register_contact_material_schedule + +.. automethod:: SimulationManager.register_particle_contact_material_schedule .. rubric:: Native point-cloud visualization diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 840c9c24c..f3dc5256c 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -22,6 +22,22 @@ decoding and compiling the constrained Expert Program schema surface. compile_mllm_expert_program decode_mllm_expert_program +embodichain.data.assets.demo_assets +----------------------------------- + +Downloadable bundles used by standalone manipulation and deformable-body +demos. Each class resolves one versioned archive into the configured +EmbodiChain data cache. + +.. currentmodule:: embodichain.data.assets.demo_assets + +.. autosummary:: + + CoordinatedPlacementAndPickment + DeformableDemoData + MultiW1Data + ScoopIceNewEnv + embodichain.data.assets.planner_assets -------------------------------------- @@ -821,15 +837,14 @@ embodichain.lab.sim.atomic_actions.transports embodichain.lab.sim.diff ------------------------ -Public differentiable-stepping bridge from manager-owned Newton trajectories -and Warp tapes into PyTorch autograd. +Public bridge from task-defined Newton kinematics and Warp tapes into PyTorch +autograd. It does not advance the Newton dynamics solver. .. currentmodule:: embodichain.lab.sim.diff .. autosummary:: NewtonStepFunc - differentiable_step tape_context embodichain.lab.sim.diff.bridge @@ -840,7 +855,6 @@ embodichain.lab.sim.diff.bridge .. autosummary:: NewtonStepFunc - differentiable_step tape_context embodichain.lab.sim.diff.runtime diff --git a/embodichain/data/assets/demo_assets.py b/embodichain/data/assets/demo_assets.py index 255002ab6..75d3637bb 100644 --- a/embodichain/data/assets/demo_assets.py +++ b/embodichain/data/assets/demo_assets.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Registered downloadable asset bundles for standalone demos.""" + from __future__ import annotations import open3d as o3d @@ -27,9 +29,23 @@ demo_assets = "demo" +__all__ = [ + "CoordinatedPlacementAndPickment", + "DeformableDemoData", + "MultiW1Data", + "ScoopIceNewEnv", +] + class ScoopIceNewEnv(EmbodiChainDataset): - def __init__(self, data_root: str = None): + """Downloadable meshes and robot assets for the scoop-ice demo.""" + + def __init__(self, data_root: str | None = None) -> None: + """Initialize the scoop-ice asset bundle. + + Args: + data_root: Optional cache root overriding the EmbodiChain default. + """ data_descriptor = o3d.data.DataDescriptor( os.path.join( EMBODICHAIN_DOWNLOAD_PREFIX, demo_assets, "ScoopIceNewEnv.zip" @@ -43,7 +59,14 @@ def __init__(self, data_root: str = None): class MultiW1Data(EmbodiChainDataset): - def __init__(self, data_root: str = None): + """Downloadable scene assets for multi-W1 manipulation demos.""" + + def __init__(self, data_root: str | None = None) -> None: + """Initialize the multi-W1 demo asset bundle. + + Args: + data_root: Optional cache root overriding the EmbodiChain default. + """ data_descriptor = o3d.data.DataDescriptor( os.path.join(EMBODICHAIN_DOWNLOAD_PREFIX, demo_assets, "multi_w1_demo.zip"), "984e8fa3aa05cb36a1fd973a475183ed", @@ -53,10 +76,37 @@ def __init__(self, data_root: str = None): super().__init__(prefix, data_descriptor, path) +class DeformableDemoData(EmbodiChainDataset): + """Shared cloth-twist and W1 T-shirt-folding demo assets.""" + + def __init__(self, data_root: str | None = None) -> None: + """Initialize the downloadable deformable-demo bundle. + + Args: + data_root: Optional cache root overriding the EmbodiChain default. + """ + data_descriptor = o3d.data.DataDescriptor( + os.path.join( + EMBODICHAIN_DOWNLOAD_PREFIX, + demo_assets, + "deformable_demo_assets.zip", + ), + "cdb1d1b105f0e96f46945052296da4d3", + ) + prefix = type(self).__name__ + path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root + super().__init__(prefix, data_descriptor, path) + + class CoordinatedPlacementAndPickment(EmbodiChainDataset): - """Dataset class for coordinated placement and pickment tutorial meshes.""" + """Downloadable meshes for coordinated placement and pickment tutorials.""" + + def __init__(self, data_root: str | None = None) -> None: + """Initialize the coordinated manipulation asset bundle. - def __init__(self, data_root: str = None): + Args: + data_root: Optional cache root overriding the EmbodiChain default. + """ data_descriptor = o3d.data.DataDescriptor( os.path.join( EMBODICHAIN_DOWNLOAD_PREFIX, diff --git a/embodichain/lab/gym/envs/differentiable_env.py b/embodichain/lab/gym/envs/differentiable_env.py index 95dfbdad0..13103e140 100644 --- a/embodichain/lab/gym/envs/differentiable_env.py +++ b/embodichain/lab/gym/envs/differentiable_env.py @@ -13,25 +13,24 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Differentiable Newton-backed EmbodiedEnv for analytic policy gradient. +"""Newton-backed kinematic environment 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. +Wraps a task-defined Newton kinematics callback in a Warp tape and bridges +autograd into PyTorch via :class:`embodichain.lab.sim.diff.NewtonStepFunc`. +The environment deliberately does not advance the configured Newton solver; +differentiable dynamics are outside the current public contract. Usage: - class MyTask(DifferentiableEmbodiedEnv): - def _apply_dynamics_action_kernel(self, action_wp, control, tape): ... + class MyTask(DifferentiableEnv): + def _apply_action_kernel(self, action_wp, tape): ... + def _make_kinematic_step_fn(self): ... def _read_outputs(self, final_state) -> dict: ... """ from __future__ import annotations -from typing import Any, Callable, Literal +from typing import Any, Callable import torch @@ -40,26 +39,24 @@ def _read_outputs(self, final_state) -> dict: ... from embodichain.lab.sim.diff import NewtonStepFunc from embodichain.utils import logger -__all__ = ["DifferentiableEmbodiedEnv"] +__all__ = ["DifferentiableEnv"] -class DifferentiableEmbodiedEnv(EmbodiedEnv): - """EmbodiedEnv variant that exposes APG-ready :py:meth:`step`. +class DifferentiableEnv(EmbodiedEnv): + """Newton-only environment with an APG-ready kinematic :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. + Subclasses implement :meth:`_apply_action_kernel`, + :meth:`_make_kinematic_step_fn`, and :meth:`_read_outputs`. The action, + kinematics, observation, and reward kernels execute inside one Warp tape. + No Newton solver or collision step is invoked by this environment. """ - differentiable_step_mode: Literal["dynamics", "kinematics"] = "dynamics" - """Stepping route used by :meth:`_build_sim_state_dict`.""" - - def __init__(self, cfg: EmbodiedEnvCfg, *args, **kwargs) -> None: + def __init__( + self, + cfg: EmbodiedEnvCfg, + *args: Any, + **kwargs: Any, + ) -> None: self._validate_diff_cfg(cfg) super().__init__(cfg, *args, **kwargs) self._truncate_backward_at: int | None = getattr( @@ -71,49 +68,25 @@ 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, " + "DifferentiableEnv requires NewtonPhysicsCfg, " f"got {type(physics_cfg).__name__}." ) if not physics_cfg.requires_grad: logger.log_error( - "DifferentiableEmbodiedEnv requires requires_grad=True on " + "DifferentiableEnv 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. + """Write an action for the task-defined kinematics callback. - 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. + The hook receives no solver control because :class:`DifferentiableEnv` + never advances Newton dynamics. """ raise NotImplementedError( - "Kinematics subclasses of DifferentiableEmbodiedEnv must implement " + "DifferentiableEnv subclasses must implement " "_apply_action_kernel(action_wp, tape)." ) @@ -127,23 +100,17 @@ def _read_outputs(self, final_state: Any) -> dict: tensors backed by ``wp.to_torch`` of grad-tracked Warp arrays. """ raise NotImplementedError( - "Subclasses of DifferentiableEmbodiedEnv must implement " - "_read_outputs(final_state)." + "DifferentiableEnv subclasses 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. + """Return the task-defined kinematics callback. Raises: - NotImplementedError: If kinematics mode has no named FK hook. + NotImplementedError: If the subclass has no kinematics hook. """ raise NotImplementedError( - "DifferentiableEmbodiedEnv in kinematics mode requires " - "_make_kinematic_step_fn()." + "DifferentiableEnv requires _make_kinematic_step_fn()." ) # -- gym surface ------------------------------------------------------ # @@ -182,66 +149,18 @@ def step(self, action: torch.Tensor): 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 = { + del action + return { "manager": self.sim, - "step_mode": mode, - "substeps": self.cfg.sim_steps_per_control, - "action_to_control_kernel": action_kernel, + "action_kernel": self._wrap_action_kernel(), "kernel_args": (), "obs_reward_fn": self._read_outputs, "last_info": {}, + "step_fn": self._make_kinematic_step_fn(), } - 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.""" + def _wrap_action_kernel(self) -> Callable[..., None]: + """Adapt the task action hook to :class:`NewtonStepFunc`.""" env = self def _inner(action_wp: Any, tape: Any, *_: Any) -> None: diff --git a/embodichain/lab/sim/_runtime_controls.py b/embodichain/lab/sim/_runtime_controls.py new file mode 100644 index 000000000..6c0b23954 --- /dev/null +++ b/embodichain/lab/sim/_runtime_controls.py @@ -0,0 +1,124 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Internal adapters for manager-owned Newton runtime controls.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from dexsim.engine.newton_physics.runtime_control import RuntimeControl + +__all__: list[str] = [] + + +class _KinematicNodalTrajectoryControl(RuntimeControl): + """Drive selected particles along offsets from their initialized positions.""" + + def __init__( + self, + target: str, + node_indices: np.ndarray, + position_offsets: np.ndarray, + *, + fps: float | None, + rebuild_self_contact_bvh: bool, + ) -> None: + self.target = target + self.node_indices = np.asarray(node_indices, dtype=np.int32).reshape(-1).copy() + self.position_offsets = ( + np.asarray(position_offsets, dtype=np.float32) + .reshape(-1, len(self.node_indices), 3) + .copy() + ) + self.fps = fps + self.rebuild_self_contact_bvh = rebuild_self_contact_bvh + self._particle_set: Any | None = None + self._initial_positions: np.ndarray | None = None + self._sample_index = 0 + self._elapsed_time = 0.0 + + def initialize(self, context: Any) -> None: + """Resolve the preconfigured inactive particles before the first substep.""" + particle_set = context.result.get_particle_set(self.target) + if np.any(self.node_indices >= particle_set.particle_count): + raise ValueError( + f"Kinematic node index exceeds particle count for {self.target!r}: " + f"max index {int(self.node_indices.max())}, particle count " + f"{particle_set.particle_count}." + ) + + positions = np.asarray( + particle_set.get_particle_positions().numpy(), + dtype=np.float32, + ).reshape(particle_set.particle_count, 3) + self._initial_positions = positions[self.node_indices].copy() + self._particle_set = particle_set + self._sample_index = 0 + self._elapsed_time = 0.0 + + def exclusive_resource_claims(self) -> tuple[object, ...]: + """Prevent multiple controls from writing the same particle set.""" + return (("kinematic_nodal_trajectory", self.target),) + + def __call__( + self, + context: Any, + substep_index: int, + substep_count: int, + substep_dt: float, + ) -> None: + """Apply the interpolated target before one Newton substep.""" + del substep_count + if self._particle_set is None or self._initial_positions is None: + raise RuntimeError("Kinematic nodal trajectory was not initialized.") + + if self.rebuild_self_contact_bvh and substep_index == 0: + rebuild_bvh = getattr(context.solver, "rebuild_bvh", None) + if callable(rebuild_bvh): + rebuild_bvh(context.current_state) + + offsets = self._current_offsets() + positions = np.asarray( + self._particle_set.get_particle_positions().numpy(), + dtype=np.float32, + ).reshape(self._particle_set.particle_count, 3) + positions[self.node_indices] = self._initial_positions + offsets + self._particle_set.set_particle_positions(positions) + + self._sample_index += 1 + self._elapsed_time += float(substep_dt) + + def _current_offsets(self) -> np.ndarray: + """Return the current sample or its time-interpolated value.""" + if self.fps is None: + sample = min(self._sample_index, len(self.position_offsets) - 1) + return self.position_offsets[sample] + + sample_position = self._elapsed_time * self.fps + lower = min(int(np.floor(sample_position)), len(self.position_offsets) - 1) + upper = min(lower + 1, len(self.position_offsets) - 1) + alpha = np.float32(sample_position - lower if upper != lower else 0.0) + return (np.float32(1.0) - alpha) * self.position_offsets[ + lower + ] + alpha * self.position_offsets[upper] + + def close(self) -> None: + """Release runtime particle references retained after initialization.""" + self._particle_set = None + self._initial_positions = None diff --git a/embodichain/lab/sim/cfg/articulation.py b/embodichain/lab/sim/cfg/articulation.py index 68eb55416..88717bef1 100644 --- a/embodichain/lab/sim/cfg/articulation.py +++ b/embodichain/lab/sim/cfg/articulation.py @@ -500,7 +500,7 @@ class ArticulationCfg(ObjectBaseCfg): 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. """ diff --git a/embodichain/lab/sim/cfg/deformable.py b/embodichain/lab/sim/cfg/deformable.py index e9a4de164..9ff3d0630 100644 --- a/embodichain/lab/sim/cfg/deformable.py +++ b/embodichain/lab/sim/cfg/deformable.py @@ -19,271 +19,128 @@ from __future__ import annotations from dataclasses import MISSING -from typing import Literal +from typing import Literal, Sequence -from dexsim.types import ( - ClothBodyAttr, - SoftBodyAttr, - SoftBodyMaterialModel, - VoxelConfig, -) +import numpy as np from embodichain.utils import configclass from ..shapes import MeshCfg from .asset import ObjectBaseCfg +__all__: list[str] = [] + @configclass class SoftbodyVoxelAttributesCfg: - # voxel config + """Newton tetrahedralization and render-volume binding parameters.""" + triangle_remesh_resolution: int = 8 - """Resolution to remesh the softbody mesh before building physics collision mesh.""" + """Resolution used to remesh the source surface before tetrahedralization.""" 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.""" + """Target face count for the proxy surface; zero disables simplification.""" simulation_mesh_resolution: int = 8 - """Resolution to build simulation voxelize textra mesh. This value must be greater than 0.""" + """Voxel resolution used to build the tetrahedral simulation mesh.""" + + voxel_num_relaxation_iters: int = 5 + """Number of tetrahedral-mesh relaxation iterations.""" + + voxel_rel_min_tet_volume: float = 0.05 + """Minimum tetrahedron volume relative to the voxel volume.""" - simulation_mesh_output_obj: bool = False - """Whether to output the simulation mesh as an obj file for debugging.""" + voxel_surface_dist_ratio: float = 0.2 + """Maximum surface distance expressed as a voxel-size ratio.""" - 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 + embedding_impl: str = "dexsim_exact_cpu" + """DexSim implementation used to bind render vertices to tetrahedra.""" @configclass class SoftbodyPhysicalAttributesCfg: - # material properties + """Newton volumetric and optional surface-element material parameters.""" + 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.""" + """Volumetric damping coefficient forwarded as Newton ``k_damp``.""" - 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.""" + density: float = 1000.0 + """Volume density in kg/m³.""" - self_collision_filter_distance: float = 0.1 - """Distance threshold below which vertex pairs may be filtered from self-collision checks.""" + surface_tri_ke: float = 0.0 + """Surface triangle elastic stiffness.""" - # --- Damping, sleep & settling --- - vertex_velocity_damping: float = 0.005 - """Per-vertex velocity damping.""" + surface_tri_ka: float = 0.0 + """Surface triangle area stiffness.""" - linear_damping: float = 0.0 - """Global linear damping applied to the soft body.""" + surface_tri_kd: float = 0.0 + """Surface triangle damping.""" - sleep_threshold: float = 0.05 - """Velocity/energy threshold below which the soft body can go to sleep.""" + surface_tri_drag: float = 0.0 + """Surface aerodynamic drag coefficient.""" - settling_threshold: float = 0.1 - """Threshold used to decide convergence/settling state.""" + surface_tri_lift: float = 0.0 + """Surface aerodynamic lift coefficient.""" - settling_damping: float = 10.0 - """Additional damping applied during settling phase.""" + add_surface_edges: bool = True + """Whether Newton creates surface bending-edge constraints.""" - # --- 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.""" + surface_edge_ke: float = 0.0 + """Surface bending-edge stiffness.""" - 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 + surface_edge_kd: float = 0.0 + """Surface bending-edge damping.""" @configclass class ClothPhysicalAttributesCfg: - # material properties - youngs: float = 1e10 - """Young's modulus (higher = stiffer).""" - - poissons: float = 0.3 - """Poisson's ratio.""" + """Newton cloth triangle, bending-edge, and spring parameters.""" - dynamic_friction: float = 0.5 - """Dynamic friction coefficient.""" - - elasticity_damping: float = 0.0 - """Elasticity damping factor.""" - - thickness: float = 0.001 - """Cloth thickness (m).""" + density: float = 1.0 + """Surface density in kg/m².""" - bending_stiffness: float = 0.00001 - """Bending stiffness.""" + tri_ke: float | None = None + """Triangle elastic stiffness; ``None`` uses the Newton default.""" - bending_damping: float = 0.0 - """Bending damping.""" + tri_ka: float | None = None + """Triangle area stiffness; ``None`` uses the Newton default.""" - # cloth body properties - enable_kinematic: bool = False - """If True, (partially) kinematic behavior is enabled.""" + tri_kd: float | None = None + """Triangle damping; ``None`` uses the Newton default.""" - enable_ccd: bool = True - """Enable continuous collision detection (CCD).""" + tri_drag: float | None = None + """Aerodynamic drag; ``None`` uses the Newton default.""" - enable_self_collision: bool = False - """Enable self-collision handling.""" + tri_lift: float | None = None + """Aerodynamic lift; ``None`` uses the Newton default.""" - has_gravity: bool = True - """Whether the cloth is affected by gravity.""" + edge_ke: float | None = None + """Bending-edge stiffness; ``None`` uses the Newton default.""" - self_collision_stress_tolerance: float = 0.9 - """Stress tolerance threshold for self-collision constraints.""" + edge_kd: float | None = None + """Bending-edge damping; ``None`` uses the Newton default.""" - collision_mesh_simplification: bool = True - """Whether to simplify the collision mesh for self-collision.""" + add_springs: bool = False + """Whether Newton creates explicit mesh-edge springs.""" - vertex_velocity_damping: float = 0.005 - """Per-vertex velocity damping.""" + spring_ke: float | None = None + """Spring stiffness; ``None`` uses the Newton default.""" - 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 + spring_kd: float | None = None + """Spring damping; ``None`` uses the Newton default.""" @configclass class DeformableObjectCfg(ObjectBaseCfg): """Common configuration contract for one deformable asset. - Concrete volume and surface configurations retain their native DexSim + Concrete volume and surface configurations author Newton particle-set properties. The discriminator is explicit so manager and visualization code do not need to infer topology from a mesh or material type. """ @@ -294,10 +151,27 @@ class DeformableObjectCfg(ObjectBaseCfg): shape: MeshCfg = MeshCfg() """Render and source-mesh configuration.""" + particle_radius: float | None = None + """Newton particle radius; ``None`` uses the active solver default.""" + + particle_flags: int | Sequence[int] | np.ndarray | None = None + """Newton particle flags, provided as one broadcast value or one value per node. + + Clear the Newton ``ACTIVE`` bit for nodes that will be driven kinematically. + Per-node arrays must follow the resolved simulation-particle order. For a + surface deformable, an array-backed + :class:`~embodichain.lab.sim.shapes.MeshCfg` preserves this order. A volume + deformable is voxelized into a separate tetrahedral simulation mesh, so its + particle indices do not correspond to source-mesh vertex indices. + """ + + validate_mesh: bool = False + """Whether Newton reports source-mesh quality validation warnings.""" + @configclass class VolumeDeformableObjectCfg(DeformableObjectCfg): - """Configuration for a volume deformable backed by DexSim ``SoftBody``.""" + """Configuration for a Newton volume-deformable particle set.""" deformable_type: Literal["volume"] = "volume" @@ -305,7 +179,7 @@ class VolumeDeformableObjectCfg(DeformableObjectCfg): """Tetrahedral simulation-mesh voxelization attributes.""" physical_attr: SoftbodyPhysicalAttributesCfg = SoftbodyPhysicalAttributesCfg() - """DexSim volume-deformable physical attributes.""" + """Newton volume-deformable physical attributes.""" @configclass @@ -315,12 +189,27 @@ class SoftObjectCfg(VolumeDeformableObjectCfg): @configclass class SurfaceDeformableObjectCfg(DeformableObjectCfg): - """Configuration for a surface deformable backed by DexSim ``ClothBody``.""" + """Configuration for a Newton surface-deformable particle set.""" deformable_type: Literal["surface"] = "surface" + visual_shape: MeshCfg | None = None + """Optional render mesh driven by the simulation surface. + + When omitted, :attr:`shape` supplies both simulation topology and rendering. + Use a separately indexed mesh here when the visual asset needs authored UVs, + seam vertices, or other detail that should not change the simulation mesh. + """ + + visual_binding_mode: Literal["auto", "nearest_vertex"] = "auto" + """Binding used to drive :attr:`visual_shape` from simulation particles. + + ``"auto"`` uses DexSim's surface embedding. ``"nearest_vertex"`` is useful + when the render mesh duplicates simulation vertices along texture seams. + """ + physical_attr: ClothPhysicalAttributesCfg = ClothPhysicalAttributesCfg() - """DexSim surface-deformable physical attributes.""" + """Newton surface-deformable physical attributes.""" @configclass diff --git a/embodichain/lab/sim/cfg/rigid_object.py b/embodichain/lab/sim/cfg/rigid_object.py index f098d7bf4..0a8fcd410 100644 --- a/embodichain/lab/sim/cfg/rigid_object.py +++ b/embodichain/lab/sim/cfg/rigid_object.py @@ -127,19 +127,19 @@ class RigidObjectGroupCfg: 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. """ diff --git a/embodichain/lab/sim/cfg/robot.py b/embodichain/lab/sim/cfg/robot.py index 505604abd..2d5647c5a 100644 --- a/embodichain/lab/sim/cfg/robot.py +++ b/embodichain/lab/sim/cfg/robot.py @@ -84,13 +84,13 @@ def _default_asset_physics_mode(self) -> AssetPhysicsMode: 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: + 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, + 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. """ diff --git a/embodichain/lab/sim/cfg/simulation.py b/embodichain/lab/sim/cfg/simulation.py index 544e03fe6..e00915058 100644 --- a/embodichain/lab/sim/cfg/simulation.py +++ b/embodichain/lab/sim/cfg/simulation.py @@ -401,6 +401,7 @@ def to_dexsim_cfg( """Convert this config to ``dexsim.engine.newton_physics.NewtonCfg``.""" from dexsim.engine.newton_physics import ( FeatherstoneSolverCfg, + MJVBDSolverCfg, MJWarpSolverCfg, NewtonCfg, NewtonCollisionPipelineCfg, @@ -420,6 +421,7 @@ def to_dexsim_cfg( solver_cfg_map = { "mujoco_warp": MJWarpSolverCfg, + "mjvbd": MJVBDSolverCfg, "xpbd": XPBDSolverCfg, "semi_implicit": SemiImplicitSolverCfg, "featherstone": FeatherstoneSolverCfg, @@ -474,6 +476,11 @@ def _normalize_newton_solver_type(solver_type: str) -> str: "mujocowarp": "mujoco_warp", "mujocowarpsolver": "mujoco_warp", "mujocowarpsolvercfg": "mujoco_warp", + "mjvbd": "mjvbd", + "mjvbdsolver": "mjvbd", + "mjvbdsolvercfg": "mjvbd", + "mjvbd_solver": "mjvbd", + "mjvbd_solver_cfg": "mjvbd", "xpbdsolver": "xpbd", "xpbdsolvercfg": "xpbd", "xpbd": "xpbd", @@ -491,8 +498,8 @@ def _normalize_newton_solver_type(solver_type: str) -> str: if key not in aliases: logger.log_error( f"Unsupported Newton solver type '{solver_type}'. " - "Expected one of 'mjwarp', 'xpbd', 'semi_implicit', " - "'featherstone', or 'vbd'." + "Expected one of 'mjwarp', 'mjvbd', 'xpbd', " + "'semi_implicit', 'featherstone', or 'vbd'." ) return aliases[key] diff --git a/embodichain/lab/sim/diff/__init__.py b/embodichain/lab/sim/diff/__init__.py index ad84e89a7..1aafcd511 100644 --- a/embodichain/lab/sim/diff/__init__.py +++ b/embodichain/lab/sim/diff/__init__.py @@ -13,24 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Differentiable Newton stepping for EmbodiChain. +"""Differentiable Newton kinematics 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. +Bridges task-defined Warp kinematics into PyTorch autograd and exposes a +:class:`tape_context` manager for advanced users composing their own kernels. +The package does not advance the Newton dynamics solver. """ 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 index 29d0fe5e0..5e0c02c25 100644 --- a/embodichain/lab/sim/diff/bridge.py +++ b/embodichain/lab/sim/diff/bridge.py @@ -13,13 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Warp-tape <-> PyTorch-autograd bridge for Newton physics.""" +"""Warp-tape ↔ PyTorch-autograd bridge for Newton kinematics.""" from __future__ import annotations from contextlib import contextmanager -import math -from typing import TYPE_CHECKING, Any, Callable, Iterator +from typing import TYPE_CHECKING, Any, Iterator import torch import warp as wp @@ -27,213 +26,83 @@ if TYPE_CHECKING: from embodichain.lab.sim.sim_manager import SimulationManager -__all__ = ["NewtonStepFunc", "differentiable_step", "tape_context"] +__all__ = ["NewtonStepFunc", "tape_context"] -def _differentiable_runtime(manager: Any) -> Any: - """Resolve Spawn's runtime while retaining lightweight test compatibility.""" +def _validate_manager(manager: Any) -> None: + """Validate the Newton gradient boundary without touching its solver.""" + if not bool(getattr(manager, "is_newton_backend", False)): + raise RuntimeError( + "Differentiable kinematics require the Newton backend with " + "requires_grad=True." + ) 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 + # Model access validates finalization and requires_grad while remaining + # independent of the configured Newton solver. + _ = runtime.model -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 _reset_tape(tape: wp.Tape | None) -> None: + """Release all arrays retained by a completed Warp tape.""" + if tape is not None: + tape.reset() -def _abort_forward( - tape: wp.Tape | None, - trajectory: Any | None, -) -> None: - """Best-effort cleanup which never masks the original forward failure.""" +def _abort_forward(tape: wp.Tape | None) -> None: + """Best-effort cleanup that never masks the original forward failure.""" try: - _reset_tape_then_release(tape, trajectory) + _reset_tape(tape) 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. + """Open a Warp tape for expert Newton kinematics kernels. + + Args: + manager: Prepared Newton-backed simulation manager in gradient mode. - Advanced users compose their own Warp kernels inside this context, then - call ``tape.backward()`` outside the with-block. + Yields: + The active Warp tape. Call ``backward`` and then ``reset`` after the + context when retaining it manually. + + Raises: + RuntimeError: If the manager does not use a finalized Newton gradient + model. """ - if not manager.is_newton_backend: - raise RuntimeError( - "tape_context requires the Newton backend with requires_grad=True." - ) + _validate_manager(manager) 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. + """Bridge one task-defined Newton kinematics step into PyTorch autograd. - 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``. + Forward records the action kernel, named kinematics callback, and output + kernels inside one Warp tape. It does not create contacts, call a Newton + solver, or advance simulation time. Backward seeds the tracked Warp output + arrays from PyTorch gradients and returns the resulting action gradient. - 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. + ``sim_state`` must contain: - 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"`` + - ``manager``: a prepared Newton-backed :class:`SimulationManager`; + - ``action_kernel``: callable ``(action_wp, tape, *kernel_args)``; + - ``kernel_args``: tuple forwarded to the action kernel; + - ``step_fn``: zero-argument task kinematics callback returning a state; + - ``obs_reward_fn``: callable that maps that state to an output dictionary. - 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`` + The output dictionary contains ``_order``, ``_grad_track``, and one torch + tensor for every name in ``_order``. ``_grad_track`` maps a name to the + backing Warp array whose gradient should be seeded, or to ``None`` for a + non-differentiable output. """ @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. - """ + """Capture ambient grad mode before PyTorch enters ``forward``.""" return super().apply(action_torch, sim_state, torch.is_grad_enabled()) @staticmethod @@ -243,21 +112,16 @@ def forward( 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"] + """Record one kinematics step and materialize its torch outputs.""" + _validate_manager(sim_state["manager"]) + action_kernel = sim_state["action_kernel"] kernel_args = sim_state["kernel_args"] + step_fn = sim_state["step_fn"] 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 - ) + if not callable(step_fn): + raise TypeError("Differentiable kinematics require a callable step_fn.") - # 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( @@ -266,55 +130,25 @@ def forward( 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) + with tape: + action_kernel(action_wp, tape, *kernel_args) + final_state = step_fn() + 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", {}) except BaseException: - _abort_forward(tape, trajectory) + _abort_forward(tape) raise if not needs_action_grad: - _reset_tape_then_release(tape, trajectory) + _reset_tape(tape) 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 @@ -326,22 +160,19 @@ def backward( ctx: Any, *grad_outputs: torch.Tensor | None, ) -> tuple[torch.Tensor | None, None, None]: + """Run Warp reverse mode and return the bridged action gradient.""" if getattr(ctx, "_bridge_released", False): raise RuntimeError( - "NewtonStepFunc backward was already consumed; create a new " - "differentiable trajectory for another backward pass." + "NewtonStepFunc backward was already consumed; run a fresh " + "kinematics step before 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( @@ -354,18 +185,13 @@ def backward( 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) + _reset_tape(ctx.tape) 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 index c6c46bdfb..abfd66b98 100644 --- a/embodichain/lab/sim/diff/runtime.py +++ b/embodichain/lab/sim/diff/runtime.py @@ -13,107 +13,33 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Differentiable transactions over a Spawn-owned Newton runtime.""" +"""Read-only access to a Spawn-owned Newton gradient model and state.""" 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. + """Expose Newton model/state buffers required by kinematic environments. - The provider is resolved for every public operation so a scene rebuild - cannot silently publish a trajectory into a replaced Newton backend. + The backend provider is resolved for every access so a scene rebuild cannot + silently return buffers owned by a replaced Spawn backend. This facade does + not expose controls, contacts, solver stepping, or gradient rollouts. """ 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." + "SimulationManager.prepare() before using differentiable " + "kinematics." ) return backend @@ -126,219 +52,33 @@ def _validated_backend(self) -> Any: ) 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." + "Differentiable Newton kinematics require requires_grad=True." ) return backend + @staticmethod + def _spawn_runtime(backend: Any) -> Any: + """Return DexSim's runtime facade across the 0.4/0.5 API boundary.""" + runtime = getattr(backend, "runtime", None) + if runtime is None: + runtime = getattr(backend, "_runtime", None) + if runtime is None: + raise RuntimeError("The Spawn-owned Newton runtime is unavailable.") + return runtime + @property def model(self) -> Any: - """Return the finalized Newton model for expert Warp operations.""" + """Return the finalized differentiable Newton model.""" 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 + backend = self._validated_backend() + return self._spawn_runtime(backend).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/deformable/__init__.py b/embodichain/lab/sim/objects/deformable/__init__.py index 91bf6b72c..eb3bbe562 100644 --- a/embodichain/lab/sim/objects/deformable/__init__.py +++ b/embodichain/lab/sim/objects/deformable/__init__.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Unified deformable-object API with DexSim volume/surface specializations.""" +"""Unified deformable-object API over Newton particle-set specializations.""" from __future__ import annotations diff --git a/embodichain/lab/sim/objects/deformable/base.py b/embodichain/lab/sim/objects/deformable/base.py index 86740a5c2..a76378cfb 100644 --- a/embodichain/lab/sim/objects/deformable/base.py +++ b/embodichain/lab/sim/objects/deformable/base.py @@ -14,15 +14,14 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Common facade for volume and surface deformable objects.""" +"""Common Newton 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 +from typing import TYPE_CHECKING, ClassVar, Literal, Sequence -import dexsim import numpy as np import torch @@ -38,22 +37,20 @@ from embodichain.utils import logger from embodichain.utils.math import matrix_from_euler, xyz_quat_to_4x4_matrix -from .data import DeformableObjectData +from .data import DeformableObjectData, _ParticleSetData if TYPE_CHECKING: - from dexsim.engine import PhysicsScene - from dexsim.spawn import SpawnResult + from dexsim.scene import Scene, SpawnedParticleSet __all__ = ["DeformableObject"] class DeformableObject(BatchEntity, ABC): - """Common facade for a batch of deformable assets. + """Common facade over a batch of Newton particle-set deformables. - 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. + Volume and surface objects retain EmbodiChain's public nodal contract, but + their runtime ownership is exclusively DexSim Spawn's Newton scene. The + Default backend and direct native soft/cloth body buffers are unsupported. """ deformable_type: ClassVar[Literal["volume", "surface"]] @@ -63,10 +60,10 @@ class DeformableObject(BatchEntity, ABC): def __init__( self, cfg: DeformableObjectCfg, - entities: Sequence[Any] | None = None, + entities: Sequence[SpawnedParticleSet] | None = None, device: torch.device = torch.device("cpu"), *, - spawn_result: SpawnResult | None = None, + spawn_result: Scene | None = None, declared_num_instances: int | None = None, ) -> None: if cfg.deformable_type != self.deformable_type: @@ -75,35 +72,40 @@ def __init__( f"{self.deformable_type!r}, got {cfg.deformable_type!r}." ) + device = torch.device(device) 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 not entities: + raise ValueError(f"A bound {type(self).__name__} requires handles.") if spawn_result is None: - self._world = dexsim.default_world() - from embodichain.lab.sim.sim_manager import get_physics_scene + raise RuntimeError( + "Deformable objects must bind to a finalized DexSim Spawn scene." + ) + if getattr(spawn_result, "backend", None) != "newton": + raise NotImplementedError( + "EmbodiChain deformable objects require the Newton backend; " + "the Default backend is no longer supported." + ) - self._ps: PhysicsScene | None = get_physics_scene() - else: - self._world = spawn_result.world - self._ps = self._world.get_physics_scene() + self._declared_num_instances = len(entities) + self._spawn_result = spawn_result + self._scene = spawn_result + self._world = spawn_result.world self._all_indices = list(range(len(entities))) + super().__init__(cfg=cfg, entities=entities, device=device) - self._data = self._create_data(entities, self._ps, device) - if spawn_result is None: - self._world.update(0.001) + self._arena_offsets = self._resolve_arena_offsets(spawn_result, entities) + self._data = self._create_data(entities, spawn_result, device) + self._local_rest_positions = self._capture_local_rest_positions() 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, @@ -111,7 +113,7 @@ def _initialize_declared( device: torch.device, declared_num_instances: int | None, ) -> None: - """Initialize a facade before Spawn materializes native handles.""" + """Initialize a facade before Spawn materializes particle handles.""" if declared_num_instances is None or declared_num_instances <= 0: raise ValueError( f"A declared {type(self).__name__} requires " @@ -120,11 +122,11 @@ def _initialize_declared( self.cfg = deepcopy(cfg) self.uid = self.cfg.uid self.device = device - self._entities: list[Any] = [] + self._entities: list[SpawnedParticleSet] = [] self._declared_num_instances = declared_num_instances self._spawn_result = None + self._scene = 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 @@ -133,25 +135,113 @@ def _initialize_declared( @abstractmethod def _create_data( self, - entities: Sequence[Any], - physics_scene: PhysicsScene, + entities: Sequence[SpawnedParticleSet], + scene: Scene, device: torch.device, - ) -> DeformableObjectData: - """Create the concrete backend data view.""" + ) -> _ParticleSetData: + """Create the topology-specific particle data view.""" + + def _initialize_topology(self, entities: Sequence[SpawnedParticleSet]) -> None: + """Capture per-instance render topology with a stable batch shape.""" + vertex_counts = tuple( + np.asarray(entity.get_render_vertices(), dtype=np.float32) + .reshape(-1, 3) + .shape[0] + for entity in entities + ) + if len(set(vertex_counts)) != 1: + raise RuntimeError( + "Replicated Newton deformable render meshes must share one " + "vertex count, but DexSim materialized counts " + f"{vertex_counts}. This indicates a render-clone topology " + "mismatch; use a compatible source mesh or one environment " + "until the DexSim clone path is corrected." + ) + triangles = [ + np.asarray(entity.get_render_triangles(), dtype=np.int32).reshape(-1, 3) + for entity in entities + ] + triangle_counts = {len(item) for item in triangles} + if len(triangle_counts) != 1: + raise ValueError( + "All instances of one deformable asset must share render " + f"triangle count, got {sorted(triangle_counts)}." + ) + for instance, (instance_triangles, vertex_count) in enumerate( + zip(triangles, vertex_counts, strict=True) + ): + if instance_triangles.size and ( + int(instance_triangles.min()) < 0 + or int(instance_triangles.max()) >= vertex_count + ): + raise ValueError( + "Deformable render topology contains an out-of-range " + f"vertex index for instance {instance}." + ) + self._surface_triangles = torch.as_tensor( + np.stack(triangles), + dtype=torch.int32, + device=self.device, + ).clone() + + @staticmethod + def _resolve_arena_offsets( + scene: Scene, + entities: Sequence[SpawnedParticleSet], + ) -> torch.Tensor: + if not scene.arenas: + offsets = np.zeros((len(entities), 3), dtype=np.float32) + else: + arena_indices = [ + scene.arenas.index(entity.arena_name) for entity in entities + ] + offsets = scene.arenas.root_offsets[arena_indices] + return torch.as_tensor(offsets, dtype=torch.float32) + + def _configured_initial_pose(self) -> torch.Tensor: + 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(4, 4) + return pose.clone() + + pose = torch.eye(4, dtype=torch.float32, device=self.device) + pose[:3, 3] = torch.as_tensor( + self.cfg.init_pos, + dtype=torch.float32, + device=self.device, + ) + rotation = ( + torch.as_tensor( + self.cfg.init_rot, + dtype=torch.float32, + device=self.device, + ) + * torch.pi + / 180.0 + ) + pose[:3, :3] = matrix_from_euler(rotation.unsqueeze(0), "XYZ")[0] + return pose - def _initialize_topology(self, entities: Sequence[Any]) -> None: - """Initialize implementation-specific surface topology.""" - del entities + def _capture_local_rest_positions(self) -> torch.Tensor: + self._require_data() + initial_pose = self._configured_initial_pose() + initial_positions = self.data.default_nodal_state_w[..., :3] + arena_offsets = self._arena_offsets.to(self.device).unsqueeze(1) + translated = initial_positions - initial_pose[:3, 3] - arena_offsets + return translated @ initial_pose[:3, :3] @property def is_spawn_bound(self) -> bool: - """Whether this facade is bound to one finalized Spawn result.""" + """Whether this facade is bound to one finalized Spawn scene.""" 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 + """Whether this facade is waiting for its Spawn scene binding.""" + return self._scene is None @property def num_instances(self) -> int: @@ -163,23 +253,31 @@ 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.""" + def attach_spawn_handles(self, entities: Sequence[SpawnedParticleSet]) -> None: + """Store materialized particle 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.""" + def bind_spawn(self, result: Scene) -> None: + """Bind a declared facade to finalized Newton particle 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, + render_body = entity.get_render_body() + project_uv = getattr(render_body, "set_projective_uv", None) + if project_uv is None: + raise NotImplementedError( + "compute_uv requires a deformable render body with " + "set_projective_uv()." + ) + project_uv(np.asarray(self.cfg.shape.project_direction)) + bound = type(self)( self.cfg, entities, self.device, spawn_result=result, ) + self.__dict__.clear() + self.__dict__.update(bound.__dict__) def __str__(self) -> str: if self.is_declared: @@ -251,31 +349,15 @@ def get_visual_material_inst( """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] - ) + """Reject legacy per-body filtering absent from Newton particle sets.""" + del filter_data, env_ids + raise NotImplementedError( + "Newton deformable collision filtering is scene/solver-owned; " + "per-object Default collision-filter data is unsupported." + ) def _resolve_env_ids(self, env_ids: Sequence[int] | None) -> list[int]: if env_ids is None: @@ -294,12 +376,10 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | None) -> list[int]: 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 - + """Set a deformable pose by transforming its captured rest particles.""" local_env_ids = self._resolve_env_ids(env_ids) if len(local_env_ids) != len(pose): - logger.log_error( + raise ValueError( f"Length of env_ids {len(local_env_ids)} does not match pose " f"length {len(pose)}." ) @@ -308,25 +388,39 @@ def set_local_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)." + raise ValueError( + f"Invalid pose shape {tuple(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.""" + """Apply rest-particle transforms through the Spawn particle batch.""" + self._require_data() + if not env_ids: + return + index = torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + local_positions = self._local_rest_positions.index_select(0, index) + rotations = pose[:, :3, :3] + translations = pose[:, :3, 3].unsqueeze(1) + arena_offsets = self._arena_offsets.to(self.device).index_select(0, index) + positions = ( + torch.bmm(local_positions, rotations.transpose(1, 2)) + + translations + + arena_offsets.unsqueeze(1) + ) + self._data._apply_nodal_state( + positions, + torch.zeros_like(positions), + env_ids, + ) def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: """Reject root-pose reads because deformables have no rigid root pose.""" @@ -336,22 +430,22 @@ def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: ) def get_current_nodal_position(self) -> torch.Tensor: - """Return current simulation-node positions in world frame.""" + """Return current simulation-particle 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.""" + """Return current simulation-particle 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]``.""" + """Return current nodal 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]``.""" + """Return the nodal state captured when Spawn was bound.""" self._require_data() return self.data.default_nodal_state_w @@ -361,53 +455,57 @@ def _require_data(self) -> None: 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.""" + """Return live render-surface vertices in world frame.""" + vertices_per_instance: list[torch.Tensor] = [] + render_pose = self._configured_initial_pose() + render_rotation = render_pose[:3, :3] + render_translation = render_pose[:3, 3] + arena_offsets = self._arena_offsets.to(self.device) + for env_idx, entity in enumerate(self._entities): + vertices_warp = entity.get_render_vertices_warp() + if vertices_warp is None: + vertices = torch.as_tensor( + entity.get_render_vertices(), + dtype=torch.float32, + device=self.device, + ).reshape(-1, 3) + else: + import warp as wp + + vertices = wp.to_torch(vertices_warp).reshape(-1, 3).to(self.device) + vertices_per_instance.append( + vertices @ render_rotation.T + + render_translation + + arena_offsets[env_idx] + ) + + vertex_counts = {len(vertices) for vertices in vertices_per_instance} + if len(vertex_counts) != 1: + raise ValueError( + "All instances of one deformable asset must share render vertex count." + ) + return torch.stack(vertices_per_instance).clone() - @abstractmethod def get_surface_triangles( self, env_ids: Sequence[int] | None = None ) -> torch.Tensor: - """Return surface triangle indices for selected environments.""" + """Return render-surface triangle indices for selected environments.""" + ids = self._resolve_env_ids(env_ids) + index = torch.as_tensor(ids, dtype=torch.long, device=self.device) + return self._surface_triangles.index_select(0, index).clone() 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.""" + """Restore the configured pose, zero 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") + initial_pose = self._configured_initial_pose() + pose = initial_pose.unsqueeze(0).repeat(len(local_env_ids), 1, 1) 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) + """Leave particle lifetime ownership with the finalized Spawn scene.""" diff --git a/embodichain/lab/sim/objects/deformable/data.py b/embodichain/lab/sim/objects/deformable/data.py index f9210e415..50077eec5 100644 --- a/embodichain/lab/sim/objects/deformable/data.py +++ b/embodichain/lab/sim/objects/deformable/data.py @@ -14,23 +14,26 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Backend-neutral data contract for deformable simulation objects.""" +"""Nodal data contract and Newton particle-set state adapter.""" from __future__ import annotations from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Sequence import torch +if TYPE_CHECKING: + from dexsim.scene import Scene, SpawnedParticleSet + __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. + Positions and velocities use the simulation world frame. Consumers can + rely on a stable ``(num_instances, num_nodes, 3)`` contract. """ @property @@ -62,3 +65,126 @@ def root_pos_w(self) -> torch.Tensor: def root_vel_w(self) -> torch.Tensor: """Return the mean nodal velocity for each deformable instance.""" return self.nodal_vel_w.mean(dim=1) + + +class _ParticleSetData(DeformableObjectData): + """Packed state adapter over DexSim 0.5 particle-set handles.""" + + def __init__( + self, + entities: Sequence[SpawnedParticleSet], + scene: Scene, + device: torch.device, + ) -> None: + self.entities = list(entities) + if not self.entities: + raise ValueError("A deformable particle-set batch cannot be empty.") + + self.scene = scene + self.device = device + self.num_instances = len(self.entities) + particle_counts = tuple(int(entity.particle_count) for entity in self.entities) + if any(count <= 0 for count in particle_counts): + raise ValueError("Deformable particle sets must contain particles.") + if len(set(particle_counts)) != 1: + raise ValueError( + "All instances of one deformable asset must have the same " + f"particle count, got {particle_counts}." + ) + + self.n_nodes = particle_counts[0] + self.batch = scene.create_particle_set_batch(self.entities) + self._position_buffer = torch.empty( + (self.num_instances, self.n_nodes, 3), + dtype=torch.float32, + device=self.device, + ) + self._velocity_buffer = torch.empty_like(self._position_buffer) + default_positions = self.nodal_pos_w + default_velocities = self.nodal_vel_w + self._default_nodal_state_w = torch.cat( + (default_positions, default_velocities), + dim=-1, + ) + + @staticmethod + def _check_batch_status(status: int | None, operation: str) -> None: + if status is not None and int(status) < 0: + raise RuntimeError( + f"DexSim particle batch failed to {operation}: status {status}." + ) + + @property + def nodal_pos_w(self) -> torch.Tensor: + """Return current Newton particle positions in world frame.""" + status = self.batch.fetch_particle_positions( + self._position_buffer.reshape(-1, 3) + ) + self._check_batch_status(status, "fetch positions") + return self._position_buffer.clone() + + @property + def nodal_vel_w(self) -> torch.Tensor: + """Return current Newton particle velocities in world frame.""" + status = self.batch.fetch_particle_velocities( + self._velocity_buffer.reshape(-1, 3) + ) + self._check_batch_status(status, "fetch velocities") + return self._velocity_buffer.clone() + + @property + def default_nodal_state_w(self) -> torch.Tensor: + """Return the particle state captured when Spawn was bound.""" + return self._default_nodal_state_w.clone() + + def _apply_nodal_state( + self, + positions: torch.Tensor, + velocities: torch.Tensor, + env_ids: Sequence[int], + ) -> None: + """Apply packed state to selected particle-set instances.""" + env_ids = [int(env_id) for env_id in env_ids] + if not env_ids: + return + if len(set(env_ids)) != len(env_ids): + raise ValueError(f"env_ids must not contain duplicates, got {env_ids}.") + + expected_shape = (len(env_ids), self.n_nodes, 3) + if tuple(positions.shape) != expected_shape: + raise ValueError( + f"positions must have shape {expected_shape}, got " + f"{tuple(positions.shape)}." + ) + if tuple(velocities.shape) != expected_shape: + raise ValueError( + f"velocities must have shape {expected_shape}, got " + f"{tuple(velocities.shape)}." + ) + + if env_ids == list(range(self.num_instances)): + batch = self.batch + else: + batch = self.scene.create_particle_set_batch( + [self.entities[env_id] for env_id in env_ids] + ) + packed_positions = ( + positions.to( + device=self.device, + dtype=torch.float32, + ) + .contiguous() + .reshape(-1, 3) + ) + packed_velocities = ( + velocities.to( + device=self.device, + dtype=torch.float32, + ) + .contiguous() + .reshape(-1, 3) + ) + position_status = batch.apply_particle_positions(packed_positions) + self._check_batch_status(position_status, "apply positions") + velocity_status = batch.apply_particle_velocities(packed_velocities) + self._check_batch_status(velocity_status, "apply velocities") diff --git a/embodichain/lab/sim/objects/deformable/surface.py b/embodichain/lab/sim/objects/deformable/surface.py index bd3df53fd..dd0498276 100644 --- a/embodichain/lab/sim/objects/deformable/surface.py +++ b/embodichain/lab/sim/objects/deformable/surface.py @@ -14,21 +14,19 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""DexSim surface-deformable object implementation.""" +"""Newton surface-deformable object implementation.""" from __future__ import annotations -from typing import Any, Sequence +from typing import TYPE_CHECKING, 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 +from .data import _ParticleSetData + +if TYPE_CHECKING: + from dexsim.scene import Scene, SpawnedClothParticleSet __all__ = [ "ClothBodyData", @@ -38,82 +36,37 @@ ] -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, - ) +class SurfaceDeformableData(_ParticleSetData): + """Newton cloth particles exposed through the legacy surface API.""" @property - def rest_vertices(self) -> torch.Tensor: - """Return rest surface vertices in simulation world frame.""" - return self._rest_position_buffer[..., :3].clone() + def particle_sets(self) -> list[SpawnedClothParticleSet]: + """Return the typed DexSim cloth particle handles.""" + return self.entities @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() + def n_vertices(self) -> int: + """Return the Newton cloth particle count per instance.""" + return self.n_nodes @property - def nodal_pos_w(self) -> torch.Tensor: - return self.vertex_position + def rest_vertices(self) -> torch.Tensor: + """Return particle positions captured when Spawn was bound.""" + return self.default_nodal_state_w[..., :3] @property - def nodal_vel_w(self) -> torch.Tensor: - return self.vertex_velocity + def vertex_position(self) -> torch.Tensor: + """Return current Newton cloth-particle positions.""" + return self.nodal_pos_w @property - def default_nodal_state_w(self) -> torch.Tensor: - return self._default_nodal_state_w.clone() + def vertex_velocity(self) -> torch.Tensor: + """Return current Newton cloth-particle velocities.""" + return self.nodal_vel_w class SurfaceDeformableObject(DeformableObject): - """A batch of DexSim surface deformables backed by ``ClothBody``.""" + """A batch of Newton cloth particle sets.""" deformable_type = "surface" spawn_kind = "cloth_object" @@ -121,116 +74,30 @@ class SurfaceDeformableObject(DeformableObject): def _create_data( self, - entities: Sequence[Any], - physics_scene: PhysicsScene, + entities: Sequence[SpawnedClothParticleSet], + scene: Scene, 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(), - ) + return SurfaceDeformableData(entities, scene, device) @property def body_data(self) -> SurfaceDeformableData | None: - """Compatibility view of the DexSim cloth data.""" + """Compatibility view of the Newton cloth particle 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.""" + """Return particle positions captured when Spawn was bound.""" self._require_data() return self.body_data.rest_vertices def get_current_vertex_position(self) -> torch.Tensor: - """Return current surface-vertex positions.""" + """Return current Newton cloth-particle positions.""" return self.get_current_nodal_position() def get_current_vertex_velocity(self) -> torch.Tensor: - """Return current surface-vertex velocities.""" + """Return current Newton cloth-particle 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 diff --git a/embodichain/lab/sim/objects/deformable/volume.py b/embodichain/lab/sim/objects/deformable/volume.py index b3ecf11ea..74f3190ea 100644 --- a/embodichain/lab/sim/objects/deformable/volume.py +++ b/embodichain/lab/sim/objects/deformable/volume.py @@ -14,24 +14,20 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""DexSim volume-deformable object implementation.""" +"""Newton volume-deformable object implementation.""" from __future__ import annotations -from functools import cached_property -from typing import Any, Sequence +from typing import TYPE_CHECKING, 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 +from .data import _ParticleSetData + +if TYPE_CHECKING: + from dexsim.scene import Scene, SpawnedSoftBodyParticleSet __all__ = [ "SoftBodyData", @@ -41,137 +37,52 @@ ] -class VolumeDeformableData(DeformableObjectData): - """DexSim soft-body buffers exposed through the common nodal contract.""" +class VolumeDeformableData(_ParticleSetData): + """Newton soft-body particles exposed through the legacy volume API.""" - 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() + @property + def particle_sets(self) -> list[SpawnedSoftBodyParticleSet]: + """Return the typed DexSim soft-body particle handles.""" + return self.entities - 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() - ) + @property + def n_collision_vertices(self) -> int: + """Return the Newton collision-particle count per instance.""" + return self.n_nodes - 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 n_sim_vertices(self) -> int: + """Return the Newton simulation-particle count per instance.""" + return self.n_nodes @property def rest_collision_vertices(self) -> torch.Tensor: - """Return rest collision vertices in simulation world frame.""" - return self._rest_position_buffer[..., :3].clone() + """Return particle positions captured when Spawn was bound.""" + return self.default_nodal_state_w[..., :3] @property def rest_sim_vertices(self) -> torch.Tensor: - """Return rest simulation vertices in simulation world frame.""" - return self._rest_sim_position_buffer[..., :3].clone() + """Return particle positions captured when Spawn was bound.""" + return self.default_nodal_state_w[..., :3] @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() + """Return current Newton collision-particle positions.""" + return self.nodal_pos_w @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() + """Return current Newton simulation-particle positions.""" + return self.nodal_pos_w @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) + """Return current Newton simulation-particle velocities.""" + return self.nodal_vel_w class VolumeDeformableObject(DeformableObject): - """A batch of DexSim volume deformables backed by ``SoftBody``.""" + """A batch of Newton volumetric soft-body particle sets.""" deformable_type = "volume" spawn_kind = "soft_object" @@ -179,102 +90,68 @@ class VolumeDeformableObject(DeformableObject): def _create_data( self, - entities: Sequence[Any], - physics_scene: PhysicsScene, + entities: Sequence[SpawnedSoftBodyParticleSet], + scene: Scene, 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 + return VolumeDeformableData(entities, scene, device) - def _apply_local_pose( + def _initialize_topology( self, - pose: torch.Tensor, - env_ids: Sequence[int], - arena_offsets: torch.Tensor, + entities: Sequence[SpawnedSoftBodyParticleSet], ) -> 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 + super()._initialize_topology(entities) + triangles = [ + np.asarray(entity.get_surface_triangles(), dtype=np.int32).reshape(-1, 3) + for entity in entities + ] + triangle_counts = {len(item) for item in triangles} + if len(triangle_counts) != 1: + raise ValueError( + "All instances of one soft body must share surface triangle " + f"count, got {sorted(triangle_counts)}." ) - sim_positions = rest_sim_local @ rotation.T + translation + arena_offset + self._collision_surface_triangles = torch.as_tensor( + np.stack(triangles), + dtype=torch.int32, + device=self.device, + ).clone() - 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) + @property + def body_data(self) -> VolumeDeformableData | None: + """Compatibility view of the Newton soft-body particle data.""" + return self._data def get_rest_collision_vertices(self) -> torch.Tensor: - """Return rest collision vertices.""" + """Return particle positions captured when Spawn was bound.""" self._require_data() return self.body_data.rest_collision_vertices def get_rest_sim_vertices(self) -> torch.Tensor: - """Return rest simulation vertices.""" + """Return particle positions captured when Spawn was bound.""" self._require_data() return self.body_data.rest_sim_vertices def get_current_collision_vertices(self) -> torch.Tensor: - """Return current collision vertices.""" + """Return current Newton collision-particle positions.""" self._require_data() return self.body_data.collision_position def get_current_sim_vertices(self) -> torch.Tensor: - """Return current simulation vertices.""" + """Return current Newton simulation-particle positions.""" return self.get_current_nodal_position() def get_current_sim_vertex_velocities(self) -> torch.Tensor: - """Return current simulation-vertex velocities.""" + """Return current Newton simulation-particle 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() + """Return the tetrahedral surface topology for selected instances.""" 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) + index = torch.as_tensor(ids, dtype=torch.long, device=self.device) + return self._collision_surface_triangles.index_select(0, index).clone() # Compatibility names retained for existing environments and tutorials. diff --git a/embodichain/lab/sim/physics/default.py b/embodichain/lab/sim/physics/default.py index c832d6a52..7bd4f71de 100644 --- a/embodichain/lab/sim/physics/default.py +++ b/embodichain/lab/sim/physics/default.py @@ -59,15 +59,13 @@ def get_scene(self): 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 + return False @property def supports_surface_deformables(self) -> bool: - return True + return False @property def supports_rigid_object_group(self) -> bool: diff --git a/embodichain/lab/sim/physics/newton.py b/embodichain/lab/sim/physics/newton.py index 0a24020aa..04180e22c 100644 --- a/embodichain/lab/sim/physics/newton.py +++ b/embodichain/lab/sim/physics/newton.py @@ -151,13 +151,11 @@ def get_scene(self): # -- capabilities --------------------------------------------------- # @property def supports_volume_deformables(self) -> bool: - # Reserved entry point: add a Newton volume adapter before enabling. - return False + return True @property def supports_surface_deformables(self) -> bool: - # Reserved entry point: add a Newton surface adapter before enabling. - return False + return True @property def supports_robot(self) -> bool: diff --git a/embodichain/lab/sim/shapes.py b/embodichain/lab/sim/shapes.py index 124edf1f8..433b54e52 100755 --- a/embodichain/lab/sim/shapes.py +++ b/embodichain/lab/sim/shapes.py @@ -16,13 +16,18 @@ from __future__ import annotations -from typing import List, Dict, Union, TYPE_CHECKING, Any from dataclasses import MISSING +from typing import Any, Dict, List, Sequence, TYPE_CHECKING, Union + +import numpy as np + from embodichain.utils import configclass, is_configclass, logger if TYPE_CHECKING: from embodichain.lab.sim.material import VisualMaterialCfg +__all__ = ["LoadOption", "ShapeCfg", "MeshCfg", "CubeCfg", "SphereCfg"] + @configclass class LoadOption: @@ -104,8 +109,27 @@ class MeshCfg(ShapeCfg): shape_type: str = "Mesh" - fpath: str = MISSING - """File path to the shape mesh file.""" + fpath: str | None = None + """File path to the shape mesh file. + + Provide either this path or both :attr:`vertices` and :attr:`triangles`. + """ + + vertices: Sequence[Sequence[float]] | np.ndarray | None = None + """Optional array-backed mesh vertices with shape ``(N, 3)``. + + Array-backed meshes preserve vertex order, which is useful when per-node + deformable flags or kinematic trajectories refer to stable node indices. + """ + + triangles: Sequence[Sequence[int]] | np.ndarray | None = None + """Optional array-backed triangle indices with shape ``(M, 3)``.""" + + normals: Sequence[Sequence[float]] | np.ndarray | None = None + """Optional per-vertex normals with shape ``(N, 3)``.""" + + uv_coords: Sequence[Sequence[float]] | np.ndarray | None = None + """Optional per-vertex texture coordinates with shape ``(N, 2)``.""" load_option: LoadOption = LoadOption() """Options for loading and processing the shape.""" diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index a2d2ea1b0..8b7484617 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -32,7 +32,16 @@ from copy import deepcopy from datetime import datetime from functools import cached_property, partial -from typing import TYPE_CHECKING, Callable, Dict, Iterator, List, Sequence, Union +from typing import ( + TYPE_CHECKING, + Callable, + Dict, + Iterator, + List, + Mapping, + Sequence, + Union, +) from dataclasses import dataclass, asdict, field, MISSING # Global cache directories @@ -179,11 +188,11 @@ def _initialize_warp_runtime(physics_cfg: PhysicsBackendCfg) -> None: 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 objects are Newton particle sets. The Default implementation is +# deliberately absent so stale soft/cloth configurations fail at declaration. _DEFORMABLE_BACKEND_IMPLEMENTATIONS = { - "default": { + "default": {}, + "newton": { "volume": ( VolumeDeformableObjectCfg, VolumeDeformableObject, @@ -197,7 +206,6 @@ def _initialize_warp_runtime(physics_cfg: PhysicsBackendCfg) -> None: "cloth_object", ), }, - "newton": {}, } @@ -1020,6 +1028,384 @@ def _init_sim_resources(self) -> None: self._default_resources = SimResources() + def register_kinematic_joint_trajectory( + self, + uid: str, + joint_positions: torch.Tensor | np.ndarray, + *, + fps: float | None = None, + root_poses: torch.Tensor | np.ndarray | None = None, + ) -> None: + """Register a Newton kinematic joint trajectory for every arena. + + The leading trajectory dimension follows EmbodiChain's batched arena + layout. Each arena row is lowered to one DexSim runtime control with + the corresponding concrete Spawn articulation path. Row zero is the + initial sample; when ``fps`` is omitted, each call to :meth:`update` + advances to the next sample. + + .. attention:: + Declare the target robot or articulation first, then call this + method before :meth:`prepare`. Runtime controls are part of the + finalized Newton simulation pipeline and cannot be added later. + + Args: + uid: UID of a robot or articulation declared on this manager. + joint_positions: Batched positions in the articulation's public + qpos order with shape ``(num_envs, frames, dof)``. + fps: Optional trajectory sample rate. When omitted, samples advance + once per EmbodiChain physics frame. + root_poses: Optional batched world-space root transforms with shape + ``(num_envs, frames, 4, 4)``. + + Raises: + RuntimeError: If the active backend is not Newton, the Spawn scene + is already finalized, or its arena count is inconsistent. + KeyError: If ``uid`` is not a declared robot or articulation. + ValueError: If an input has an invalid batch shape or non-finite + values. + """ + if not isinstance(uid, str) or not uid: + raise ValueError("uid must be a non-empty string.") + if not self.is_newton_backend: + raise RuntimeError( + "Kinematic joint trajectory controls require the Newton backend." + ) + if uid not in self._robots and uid not in self._articulations: + raise KeyError(f"Robot or articulation {uid!r} is not declared.") + + scene = self._spawn_scene + if scene.builder.is_finalized: + raise RuntimeError( + "Kinematic joint trajectories must be registered before " + "SimulationManager.prepare()." + ) + + if isinstance(joint_positions, torch.Tensor): + positions = joint_positions.detach().cpu().numpy() + else: + positions = np.asarray(joint_positions) + positions = np.asarray(positions, dtype=np.float32) + expected_prefix = (self.num_envs,) + if ( + positions.ndim != 3 + or positions.shape[:1] != expected_prefix + or positions.shape[1] == 0 + or positions.shape[2] == 0 + ): + raise ValueError( + "joint_positions must have non-empty shape " + f"({self.num_envs}, frames, dof); got {positions.shape}." + ) + if not np.isfinite(positions).all(): + raise ValueError("joint_positions must contain only finite values.") + + poses: np.ndarray | None = None + if root_poses is not None: + if isinstance(root_poses, torch.Tensor): + poses = root_poses.detach().cpu().numpy() + else: + poses = np.asarray(root_poses) + poses = np.asarray(poses, dtype=np.float32) + expected_shape = (self.num_envs, positions.shape[1], 4, 4) + if poses.shape != expected_shape: + raise ValueError( + f"root_poses must have shape {expected_shape}; got {poses.shape}." + ) + if not np.isfinite(poses).all(): + raise ValueError("root_poses must contain only finite values.") + + arena_names = scene.arena_names + if len(arena_names) != self.num_envs: + raise RuntimeError( + "Spawn arena count does not match SimulationManager.num_envs: " + f"{len(arena_names)} != {self.num_envs}." + ) + + from dexsim.engine.newton_physics import KinematicJointTrajectoryControl + + controls = tuple( + KinematicJointTrajectoryControl( + f"{arena_name}/{uid}", + positions[env_index], + fps=fps, + root_poses=None if poses is None else poses[env_index], + ) + for env_index, arena_name in enumerate(arena_names) + ) + for control in controls: + scene.builder.add_runtime_control(control) + + def register_contact_material_schedule( + self, + uid: str, + keyframes: Mapping[str, Sequence[Sequence[float]]], + *, + link_names: Sequence[str] | None = None, + ) -> None: + """Register time-varying Newton contact properties for an asset. + + The manager expands the declared UID to the concrete Spawn path in + every Arena. Supported keyframe tracks are ``dynamic_friction``, + ``stiffness``, and ``damping``; each track contains ``(time, value)`` + pairs and is sampled piecewise-constantly in simulation time. + + .. attention:: + Declare the rigid object, robot, or articulation first and call + this method before :meth:`prepare`. Runtime controls are part of + the finalized Newton pipeline and cannot be added later. + + Args: + uid: UID of a declared rigid object, robot, or articulation. + keyframes: Contact-property tracks keyed by property name. + link_names: Optional articulation-link names to update. ``None`` + updates every collision shape belonging to the target. + + Raises: + RuntimeError: If the backend is not Newton, the scene is already + finalized, or its Arena count is inconsistent. + KeyError: If ``uid`` does not identify a declared supported asset. + ValueError: If the UID or keyframes are invalid. + """ + if not isinstance(uid, str) or not uid: + raise ValueError("uid must be a non-empty string.") + if not self.is_newton_backend: + raise RuntimeError("Contact material schedules require the Newton backend.") + if ( + uid not in self._rigid_objects + and uid not in self._robots + and uid not in self._articulations + ): + raise KeyError( + f"Rigid object, robot, or articulation {uid!r} is not declared." + ) + + scene = self._spawn_scene + if scene.builder.is_finalized: + raise RuntimeError( + "Contact material schedules must be registered before " + "SimulationManager.prepare()." + ) + arena_names = scene.arena_names + if len(arena_names) != self.num_envs: + raise RuntimeError( + "Spawn arena count does not match SimulationManager.num_envs: " + f"{len(arena_names)} != {self.num_envs}." + ) + + from dexsim.engine.newton_physics import ContactMaterialSchedule + + controls = tuple( + ContactMaterialSchedule( + f"{arena_name}/{uid}", + keyframes, + link_names=link_names, + ) + for arena_name in arena_names + ) + for control in controls: + scene.builder.add_runtime_control(control) + + def register_particle_contact_material_schedule( + self, + keyframes: Mapping[str, Sequence[Sequence[float]]], + ) -> None: + """Register time-varying Newton particle contact properties. + + Supported tracks are ``dynamic_friction``, ``stiffness``, and + ``damping``. Values apply scene-wide to particle-versus-rigid contacts + and are sampled piecewise-constantly in simulation time. + + .. attention:: + Register this control before :meth:`prepare`. This host-side + control requires direct Newton stepping and therefore disables + CUDA Graph replay for the finalized simulation. + + Args: + keyframes: Particle contact-property tracks containing ``(time, + value)`` pairs. + + Raises: + RuntimeError: If the backend is not Newton or the scene is already + finalized. + ValueError: If the keyframes are invalid. + """ + if not self.is_newton_backend: + raise RuntimeError( + "Particle contact material schedules require the Newton backend." + ) + scene = self._spawn_scene + if scene.builder.is_finalized: + raise RuntimeError( + "Particle contact material schedules must be registered before " + "SimulationManager.prepare()." + ) + + from dexsim.engine.newton_physics import ParticleContactMaterialSchedule + + scene.builder.add_runtime_control(ParticleContactMaterialSchedule(keyframes)) + + def register_kinematic_nodal_trajectory( + self, + uid: str, + node_indices: torch.Tensor | np.ndarray | Sequence[int], + position_offsets: torch.Tensor | np.ndarray, + *, + fps: float | None = None, + rebuild_self_contact_bvh: bool = False, + ) -> None: + """Register a Newton trajectory for selected deformable nodes. + + Each selected node is fixed during Newton model construction and then + moved relative to the world position captured when the runtime control + initializes. The manager expands the batched offsets to one control per + Arena without exposing the private Spawn scene. + + When ``fps`` is provided, samples are linearly interpolated at Newton + substep times. Otherwise one sample is consumed per substep. The final + sample is held after the trajectory ends. + + .. attention:: + Declare the deformable first and clear the Newton ``ACTIVE`` bit in + its ``particle_flags`` for every selected node. Then register this + control before :meth:`prepare`. This host-side control makes Newton + use direct substep launches instead of CUDA Graph replay. Surface + node indices can follow an array-backed mesh directly; volume node + indices refer to the generated tetrahedral simulation particles, + not source-mesh vertices. + + Args: + uid: UID of a deformable declared on this manager. + node_indices: Shared one-dimensional simulation-particle indices. + position_offsets: Batched world-frame position offsets with shape + ``(num_envs, samples, selected_nodes, 3)``. + fps: Optional trajectory sample rate in samples per second. + rebuild_self_contact_bvh: Whether to request a full solver BVH + rebuild at the start of every physics frame when supported. + + Raises: + RuntimeError: If the active backend is not Newton, the Spawn scene + is already finalized, or its Arena count is inconsistent. + KeyError: If ``uid`` is not a declared deformable. + TypeError: If node indices, ``fps``, or the BVH option have invalid + types. + ValueError: If an input has an invalid shape or value, or selected + nodes were not configured as inactive particles. + """ + if not isinstance(uid, str) or not uid: + raise ValueError("uid must be a non-empty string.") + if not self.is_newton_backend: + raise RuntimeError( + "Kinematic nodal trajectory controls require the Newton backend." + ) + if uid not in self._deformable_objects: + raise KeyError(f"Deformable object {uid!r} is not declared.") + + scene = self._spawn_scene + if scene.builder.is_finalized: + raise RuntimeError( + "Kinematic nodal trajectories must be registered before " + "SimulationManager.prepare()." + ) + + if isinstance(node_indices, torch.Tensor): + raw_indices = node_indices.detach().cpu().numpy() + else: + raw_indices = np.asarray(node_indices) + if raw_indices.ndim != 1 or raw_indices.size == 0: + raise ValueError("node_indices must be a non-empty one-dimensional array.") + if raw_indices.dtype.kind not in "iu": + raise TypeError("node_indices must contain integers.") + if np.any(raw_indices < 0): + raise ValueError("node_indices must be non-negative.") + if np.any(raw_indices > np.iinfo(np.int32).max): + raise ValueError("node_indices exceed the supported int32 range.") + indices = np.asarray(raw_indices, dtype=np.int32) + if len(np.unique(indices)) != len(indices): + raise ValueError("node_indices must not contain duplicates.") + + configured_flags = self._deformable_objects[uid].cfg.particle_flags + if configured_flags is None: + raise ValueError( + "Kinematic nodes require particle_flags with the Newton ACTIVE " + "bit cleared before model construction." + ) + flags = np.asarray(configured_flags) + if flags.ndim == 0: + selected_flags = np.full(len(indices), int(flags), dtype=np.int64) + elif flags.ndim == 1: + if int(indices.max()) >= len(flags): + raise ValueError( + "node_indices exceed the configured particle_flags length: " + f"max index {int(indices.max())}, length {len(flags)}." + ) + selected_flags = flags[indices] + else: + raise ValueError( + "Configured particle_flags must be scalar or one-dimensional." + ) + newton_active_particle_flag = 1 + if np.any( + np.asarray(selected_flags, dtype=np.int64) & newton_active_particle_flag + ): + raise ValueError( + "Every kinematic node must have the Newton ACTIVE particle flag cleared." + ) + + if isinstance(position_offsets, torch.Tensor): + offsets = position_offsets.detach().cpu().numpy() + else: + offsets = np.asarray(position_offsets) + offsets = np.asarray(offsets, dtype=np.float32) + if ( + offsets.ndim != 4 + or offsets.shape[0] != self.num_envs + or offsets.shape[1] == 0 + or offsets.shape[2:] != (len(indices), 3) + ): + raise ValueError( + "position_offsets must have non-empty shape " + f"({self.num_envs}, samples, {len(indices)}, 3); got " + f"{offsets.shape}." + ) + if not np.isfinite(offsets).all(): + raise ValueError("position_offsets must contain only finite values.") + + if fps is not None: + if isinstance(fps, bool) or not isinstance( + fps, (int, float, np.integer, np.floating) + ): + raise TypeError("fps must be a finite positive number or None.") + fps = float(fps) + if not np.isfinite(fps) or fps <= 0.0: + raise ValueError("fps must be a finite positive number.") + if not isinstance(rebuild_self_contact_bvh, bool): + raise TypeError("rebuild_self_contact_bvh must be a bool.") + + arena_names = scene.arena_names + if len(arena_names) != self.num_envs: + raise RuntimeError( + "Spawn arena count does not match SimulationManager.num_envs: " + f"{len(arena_names)} != {self.num_envs}." + ) + + from embodichain.lab.sim._runtime_controls import ( + _KinematicNodalTrajectoryControl, + ) + + controls = tuple( + _KinematicNodalTrajectoryControl( + f"{arena_name}/{uid}", + indices, + offsets[env_index], + fps=fps, + rebuild_self_contact_bvh=rebuild_self_contact_bvh, + ) + for env_index, arena_name in enumerate(arena_names) + ) + for control in controls: + scene.builder.add_runtime_control(control) + def prepare(self) -> None: """Materialize declarations, bind state, and resolve sensor parents.""" scene = self._spawn_scene @@ -1111,55 +1497,6 @@ def finalize_newton_physics(self) -> None: """ 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 self.differentiable_runtime.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`. - - 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``. - - 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. @@ -1881,10 +2218,8 @@ def add_rigid_object( 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. + Deformables are DexSim 0.5 typed particle sets owned by the Newton + Spawn scene. The Default backend is intentionally unsupported. Args: cfg: Volume- or surface-deformable configuration. @@ -1909,12 +2244,26 @@ def add_deformable_object(self, cfg: DeformableObjectCfg) -> DeformableObject: ) if not supported: raise NotImplementedError( - f"The {self.physics.name} backend does not yet provide a " - f"{deformable_type}-deformable object adapter." + "EmbodiChain deformable objects require the Newton backend; " + f"the {self.physics.name} backend does not support them." ) if self.device.type != "cuda": raise NotImplementedError( - "DexSim deformable objects currently require a CUDA device." + "Newton deformable particle sets currently require a CUDA device." + ) + solver_type = self._active_newton_solver_type + supported_solvers = {"xpbd", "semi_implicit", "vbd", "mjvbd"} + if solver_type not in supported_solvers: + raise NotImplementedError( + f"Newton solver {solver_type!r} does not support deformable " + "particle sets; select one of 'xpbd', 'semi_implicit', 'vbd', " + "or 'mjvbd'." + ) + physics_cfg = self.sim_config.physics_cfg + if isinstance(physics_cfg, NewtonPhysicsCfg) and physics_cfg.requires_grad: + raise NotImplementedError( + "Newton deformable state mutation is unavailable when " + "requires_grad=True." ) if self.spawn_result is not None: raise NotImplementedError( diff --git a/embodichain/lab/sim/spawn/descriptors.py b/embodichain/lab/sim/spawn/descriptors.py index bf92ce2ce..7185cf0da 100644 --- a/embodichain/lab/sim/spawn/descriptors.py +++ b/embodichain/lab/sim/spawn/descriptors.py @@ -45,11 +45,9 @@ ClothPhysicsDesc, CollisionApproximation, CollisionDesc, - DexsimClothPhysicsDesc, DexsimCollisionDesc, DexsimJointDesc, DexsimPhysicsDesc, - DexsimSoftBodyPhysicsDesc, GeometryDesc, MaterialDesc, NewtonCollisionDesc, @@ -439,7 +437,11 @@ def rigid_desc_from_cfg( ) -> 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): + if ( + isinstance(cfg.shape, MeshCfg) + and not _is_missing(cfg.shape.fpath) + 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." @@ -491,24 +493,139 @@ def rigid_desc_from_cfg( return descriptor, materials +def _particle_flags_from_cfg( + flags: int | Sequence[int] | np.ndarray | None, +) -> int | np.ndarray | None: + """Validate and copy optional Newton particle flags.""" + max_flag = int(np.iinfo(np.int32).max) + if flags is None: + return None + if np.isscalar(flags): + if isinstance(flags, (bool, np.bool_)) or not isinstance( + flags, numbers.Integral + ): + raise TypeError("particle_flags scalar must be an integer bitmask.") + value = int(flags) + if value < 0 or value > max_flag: + raise ValueError(f"particle_flags values must lie in [0, {max_flag}].") + return value + + values = np.asarray(flags) + if values.ndim != 1 or values.size == 0: + raise ValueError("particle_flags must be a non-empty one-dimensional array.") + if values.dtype.kind not in "iu": + raise TypeError("particle_flags array must contain integer bitmasks.") + if np.any(values < 0) or np.any(values > max_flag): + raise ValueError(f"particle_flags values must lie in [0, {max_flag}].") + return values.astype(np.int32, copy=True) + + +def _mesh_geometry_from_cfg(shape: MeshCfg, *, segment_name: str) -> GeometryDesc: + """Compile one file-backed or array-backed mesh configuration.""" + has_file = ( + not _is_missing(shape.fpath) + and shape.fpath is not None + and bool(str(shape.fpath).strip()) + ) + has_vertices = shape.vertices is not None + has_triangles = shape.triangles is not None + + if has_vertices != has_triangles: + raise ValueError( + "MeshCfg.vertices and MeshCfg.triangles must be provided together." + ) + if has_file and has_vertices: + raise ValueError( + "MeshCfg must provide either fpath or vertices/triangles, not both." + ) + if not has_file and not has_vertices: + raise ValueError( + "MeshCfg must provide a non-empty fpath or vertices/triangles." + ) + if has_file: + if shape.normals is not None or shape.uv_coords is not None: + raise ValueError( + "MeshCfg.normals and uv_coords require an array-backed mesh." + ) + return GeometryDesc.mesh(file_path=str(shape.fpath), segment_name=segment_name) + + vertices = np.asarray(shape.vertices, dtype=np.float32) + if vertices.ndim != 2 or vertices.shape[1:] != (3,) or len(vertices) == 0: + raise ValueError("MeshCfg.vertices must have non-empty shape (N, 3).") + if not np.isfinite(vertices).all(): + raise ValueError("MeshCfg.vertices must contain only finite values.") + + raw_triangles = np.asarray(shape.triangles) + if raw_triangles.dtype.kind not in "iu": + raise TypeError("MeshCfg.triangles must contain integer indices.") + if ( + raw_triangles.ndim != 2 + or raw_triangles.shape[1:] != (3,) + or len(raw_triangles) == 0 + ): + raise ValueError("MeshCfg.triangles must have non-empty shape (M, 3).") + if np.any(raw_triangles < 0) or np.any(raw_triangles >= len(vertices)): + raise ValueError("MeshCfg.triangles contain an out-of-range vertex index.") + triangles = raw_triangles.astype(np.int32, copy=True) + + normals = None + if shape.normals is not None: + normals = np.asarray(shape.normals, dtype=np.float32) + if normals.shape != vertices.shape or not np.isfinite(normals).all(): + raise ValueError( + f"MeshCfg.normals must have finite shape {vertices.shape}." + ) + normals = normals.copy() + + uv_coords = None + if shape.uv_coords is not None: + uv_coords = np.asarray(shape.uv_coords, dtype=np.float32) + if uv_coords.shape != (len(vertices), 2) or not np.isfinite(uv_coords).all(): + raise ValueError( + f"MeshCfg.uv_coords must have finite shape ({len(vertices)}, 2)." + ) + uv_coords = uv_coords.copy() + + return GeometryDesc.mesh( + vertices=vertices.copy(), + triangles=triangles, + normals=normals, + uv_coords=uv_coords, + segment_name=segment_name, + ) + + 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.""" + """Translate a volume deformable into a Newton particle-set 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) + geometry = _mesh_geometry_from_cfg(cfg.shape, 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) + density = float(physical_attr.density) + particle_radius = ( + None if cfg.particle_radius is None else float(cfg.particle_radius) + ) + if not math.isfinite(youngs) or youngs < 0.0: + raise ValueError("Soft-body youngs must be a finite non-negative value.") + if not math.isfinite(poissons) or not -1.0 < poissons < 0.5: + raise ValueError("Soft-body poissons must be finite and lie in (-1, 0.5).") + if not math.isfinite(density) or density <= 0.0: + raise ValueError("Soft-body density must be a finite positive value.") + if particle_radius is not None and ( + not math.isfinite(particle_radius) or particle_radius <= 0.0 + ): + raise ValueError( + "Soft-body particle_radius must be finite and positive when set." + ) + descriptor = SoftBodyDesc( name=uid, pose=_pose_from_cfg(cfg), @@ -518,18 +635,31 @@ def volume_deformable_desc_from_cfg( material_ref=material_ref, ), physics=SoftBodyPhysicsDesc( - volume_density=float(physical_attr.density), + volume_density=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)), + k_damp=float(physical_attr.elasticity_damping), + surface_tri_ke=float(physical_attr.surface_tri_ke), + surface_tri_ka=float(physical_attr.surface_tri_ka), + surface_tri_kd=float(physical_attr.surface_tri_kd), + surface_tri_drag=float(physical_attr.surface_tri_drag), + surface_tri_lift=float(physical_attr.surface_tri_lift), + add_surface_edges=bool(physical_attr.add_surface_edges), + surface_edge_ke=float(physical_attr.surface_edge_ke), + surface_edge_kd=float(physical_attr.surface_edge_kd), ), - # 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, + voxel_num_relaxation_iters=cfg.voxel_attr.voxel_num_relaxation_iters, + voxel_rel_min_tet_volume=cfg.voxel_attr.voxel_rel_min_tet_volume, + voxel_surface_dist_ratio=cfg.voxel_attr.voxel_surface_dist_ratio, + embedding_impl=cfg.voxel_attr.embedding_impl, ), + particle_flags=_particle_flags_from_cfg(cfg.particle_flags), + particle_radius=particle_radius, + validate_mesh=cfg.validate_mesh, per_env=per_env, ) materials = {} if material_entry is None else {material_entry[0]: material_entry[1]} @@ -541,28 +671,75 @@ def surface_deformable_desc_from_cfg( *, per_env: bool = True, ) -> tuple[ClothDesc, dict[str, MaterialDesc]]: - """Translate a surface-deformable config into a DexSim descriptor.""" + """Translate a surface deformable into a Newton particle-set descriptor.""" uid = _required_uid(cfg.uid, "Surface deformable") - if _is_missing(cfg.shape.fpath) or not str(cfg.shape.fpath).strip(): + geometry = _mesh_geometry_from_cfg(cfg.shape, segment_name=uid) + if cfg.visual_binding_mode not in {"auto", "nearest_vertex"}: raise ValueError( - "SurfaceDeformableObjectCfg.shape.fpath must be a non-empty path." + "Surface deformable visual_binding_mode must be 'auto' or " + f"'nearest_vertex'; got {cfg.visual_binding_mode!r}." + ) + render_shape = cfg.shape if cfg.visual_shape is None else cfg.visual_shape + visual_geometry = ( + None + if cfg.visual_shape is None + else _mesh_geometry_from_cfg( + cfg.visual_shape, + segment_name=f"{uid}_visual", ) - geometry = GeometryDesc.mesh(file_path=str(cfg.shape.fpath), segment_name=uid) + ) material_ref, material_entry = _compile_visual_material( - uid, cfg.shape.visual_material + uid, render_shape.visual_material ) + physical_attr = cfg.physical_attr + density = float(physical_attr.density) + particle_radius = ( + None if cfg.particle_radius is None else float(cfg.particle_radius) + ) + if not math.isfinite(density) or density <= 0.0: + raise ValueError("Cloth density must be a finite positive value.") + if particle_radius is not None and ( + not math.isfinite(particle_radius) or particle_radius <= 0.0 + ): + raise ValueError("Cloth particle_radius must be finite and positive when set.") 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, + mesh=( + RenderDesc.from_geometry( + geometry, + load_option=_compile_load_option(cfg.shape), + material_ref=material_ref, + ) + if visual_geometry is None + else geometry + ), + visual_mesh=( + None + if visual_geometry is None + else RenderDesc.from_geometry( + visual_geometry, + load_option=_compile_load_option(render_shape), + material_ref=material_ref, + ) ), + visual_binding_mode=cfg.visual_binding_mode, physics=ClothPhysicsDesc( - surface_density=float(cfg.physical_attr.density), - dexsim=DexsimClothPhysicsDesc(**_configured_values(cfg.physical_attr)), + surface_density=density, + tri_ke=physical_attr.tri_ke, + tri_ka=physical_attr.tri_ka, + tri_kd=physical_attr.tri_kd, + tri_drag=physical_attr.tri_drag, + tri_lift=physical_attr.tri_lift, + edge_ke=physical_attr.edge_ke, + edge_kd=physical_attr.edge_kd, + add_springs=bool(physical_attr.add_springs), + spring_ke=physical_attr.spring_ke, + spring_kd=physical_attr.spring_kd, ), + particle_flags=_particle_flags_from_cfg(cfg.particle_flags), + particle_radius=particle_radius, + validate_mesh=cfg.validate_mesh, per_env=per_env, ) materials = {} if material_entry is None else {material_entry[0]: material_entry[1]} @@ -1398,8 +1575,7 @@ def _compile_geometry( ) -> 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.") + geometry = _mesh_geometry_from_cfg(shape, segment_name=cfg.uid or "mesh") max_hulls, acd_method, sdf_resolution = _resolved_mesh_collision_settings( cfg, physics=physics, @@ -1428,9 +1604,7 @@ def _compile_geometry( "its cooking resolution." ) return ( - GeometryDesc.mesh( - file_path=str(shape.fpath), segment_name=cfg.uid or "mesh" - ), + geometry, approximation, max(1, max_hulls), ) diff --git a/embodichain/lab/sim/utility/sim_utils.py b/embodichain/lab/sim/utility/sim_utils.py index ab7f81303..1b54d6b28 100644 --- a/embodichain/lab/sim/utility/sim_utils.py +++ b/embodichain/lab/sim/utility/sim_utils.py @@ -748,58 +748,20 @@ def load_mesh_objects_from_cfg( def load_soft_object_from_cfg( cfg: SoftObjectCfg, env_list: List[Arena] ) -> List[MeshObject]: - obj_list = [] - - 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 - option.share_mesh = False - - for i, env in enumerate(env_list): - obj = env.load_actor( - fpath=cfg.shape.fpath, duplicate=True, attach_scene=True, option=option - ) - obj.add_softbody(cfg.voxel_attr.attr(), cfg.physical_attr.attr()) - if cfg.shape.compute_uv: - vertices = obj.get_vertices() - triangles = obj.get_triangles() - - o3d_mesh = o3d.t.geometry.TriangleMesh(vertices, triangles) - _, uvs = get_mesh_auto_uv(o3d_mesh, cfg.shape.project_direction) - obj.set_uv_mapping(uvs) - obj.set_name(f"{cfg.uid}_{i}") - obj_list.append(obj) - return obj_list + """Reject the removed direct-native soft-body construction path.""" + del cfg, env_list + raise NotImplementedError( + "Direct soft-body loading was removed. Configure Newton and call " + "SimulationManager.add_soft_object() before prepare()." + ) def load_cloth_object_from_cfg( cfg: ClothObjectCfg, env_list: List[Arena] ) -> List[MeshObject]: - obj_list = [] - - 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 - option.share_mesh = False - - for i, env in enumerate(env_list): - obj = env.load_actor( - fpath=cfg.shape.fpath, duplicate=True, attach_scene=True, option=option - ) - obj.add_clothbody(cfg.physical_attr.attr()) - if cfg.shape.compute_uv: - vertices = obj.get_vertices() - triangles = obj.get_triangles() - - o3d_mesh = o3d.t.geometry.TriangleMesh(vertices, triangles) - _, uvs = get_mesh_auto_uv(o3d_mesh, cfg.shape.project_direction) - obj.set_uv_mapping(uvs) - obj.set_name(f"{cfg.uid}_{i}") - obj_list.append(obj) - return obj_list + """Reject the removed direct-native cloth construction path.""" + del cfg, env_list + raise NotImplementedError( + "Direct cloth loading was removed. Configure Newton and call " + "SimulationManager.add_cloth_object() before prepare()." + ) diff --git a/embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py b/embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py index 717e5dc53..7a519ae17 100644 --- a/embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py +++ b/embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py @@ -13,14 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Franka FR3 reach task with differentiable Newton physics (APG). +"""Franka FR3 reach task with differentiable Newton kinematics (APG). -Built on :class:`DifferentiableEmbodiedEnv`. The Warp-tape bridge +Built on :class:`DifferentiableEnv`. 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 +forward-kinematics path (``newton.eval_fk``). The configured semi-implicit +solver is not advanced; the task runs FK directly, matching the reference APG implementation in ``/root/sources/analytic_policy_gradients/envs/franka_reach_env.py``. """ @@ -35,7 +33,7 @@ import newton import newton.utils -from embodichain.lab.gym.envs.differentiable_env import DifferentiableEmbodiedEnv +from embodichain.lab.gym.envs.differentiable_env import DifferentiableEnv from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg from embodichain.lab.gym.utils.registration import register_env from embodichain.lab.sim.cfg import ( @@ -102,7 +100,7 @@ def _reach_reward_kernel( @register_env("FrankaReachApg-v0") -class FrankaReachApgEnv(DifferentiableEmbodiedEnv): +class FrankaReachApgEnv(DifferentiableEnv): """Differentiable Franka FR3 reach task for analytic policy gradients. The environment resolves the Franka FR3 URDF via @@ -116,16 +114,12 @@ class FrankaReachApgEnv(DifferentiableEmbodiedEnv): 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. + The configured semi-implicit solver is part of Newton scene setup but is + never advanced. This matches the current kinematics-only + :class:`DifferentiableEnv` contract. """ metadata = {"render_modes": ["human"], "default_num_envs": 4} - differentiable_step_mode = "kinematics" def __init__( self, @@ -278,7 +272,7 @@ def _sample_new_targets(self, env_ids: torch.Tensor) -> None: n, -1 ) - # -- DifferentiableEmbodiedEnv contract ------------------------------ # + # -- DifferentiableEnv contract -------------------------------------- # def _build_sim_state_dict(self, action: torch.Tensor) -> dict: """Detach FK primal buffers before the parent opens a Warp tape.""" @@ -290,10 +284,8 @@ def _build_sim_state_dict(self, action: torch.Tensor) -> dict: 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 + The environment intentionally bypasses its configured solver and runs + forward kinematics directly inside the tape. ``self._new_joint_q`` is populated by :meth:`_apply_action_kernel` before this callable runs. """ env = self @@ -413,7 +405,7 @@ def _read_outputs(self, final_state: Any) -> dict: 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 + The parent :meth:`DifferentiableEnv.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 diff --git a/examples/sim/demo/cloth_twist.py b/examples/sim/demo/cloth_twist.py new file mode 100644 index 000000000..b3fc48f28 --- /dev/null +++ b/examples/sim/demo/cloth_twist.py @@ -0,0 +1,376 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Twist two fixed cloth edges in opposite directions with Newton VBD.""" + +from __future__ import annotations + +import argparse +import math +from pathlib import Path + +import numpy as np + +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 ( + ClothObjectCfg, + ClothPhysicalAttributesCfg, + NewtonPhysicsCfg, + RenderCfg, +) +from embodichain.lab.sim.material import VisualMaterialCfg +from embodichain.lab.sim.objects import ClothObject +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.visualization import visualization_cfg_from_args +from embodichain.utils import logger + +ASSET_DATASET = "DeformableDemoData" +TEXTURE_FILE_MAP = { + "mianbu": "mianbu.png", + "shabu": "shabu.png", + "mabu": "mabu.png", + "pige": "pige.png", + "jinduan": "jinduan.png", + "niuzai": "niuzai.png", +} + +FPS = 60 +NUM_SUBSTEPS = 10 +SOLVER_ITERATIONS = 4 +DEFAULT_FRAMES = 1000 +ROTATION_ANGULAR_VELOCITY = math.pi / 3.0 +ROTATION_END_TIME = 30.0 +MESH_SCALE = 0.01 +GRID_SIZE = 50 +CLOTH_POSITION = (0.0, 0.0, 0.75) + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the cloth-twist demo.""" + parser = argparse.ArgumentParser(description=__doc__) + add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--iterations", + type=int, + default=DEFAULT_FRAMES, + help="Number of 60 Hz simulation frames.", + ) + parser.add_argument( + "--cloth-material", + choices=tuple(TEXTURE_FILE_MAP), + default="jinduan", + help="Texture preset packaged with the DexSim reference demo.", + ) + parser.set_defaults(device="cuda", physics="newton") + args = parser.parse_args() + if args.physics != "newton": + parser.error("Cloth requires --physics newton.") + if not str(args.device).startswith("cuda"): + parser.error("EmbodiChain cloth currently requires a CUDA device.") + if args.num_envs != 1: + parser.error("This cloth demo currently supports --num_envs 1.") + if args.iterations <= 0: + parser.error("--iterations must be positive.") + return args + + +def prepare_cloth_asset( + material_name: str, +) -> tuple[Path, np.ndarray, np.ndarray, np.ndarray]: + """Resolve the reference assets and load an order-preserving mesh. + + Args: + material_name: Texture preset selected on the command line. + + Returns: + The texture path, scaled vertices, triangles, and per-vertex UVs. + + Raises: + FileNotFoundError: If the downloaded reference package is incomplete. + RuntimeError: If the source mesh does not have the expected topology. + """ + asset_root = Path(get_data_path(f"{ASSET_DATASET}/cloth_twist")).parent + source_mesh = asset_root / "cloth_twist" / "cloth_twist_square_cloth.obj" + texture_path = asset_root / "textures" / TEXTURE_FILE_MAP[material_name] + + missing = [path for path in (source_mesh, texture_path) if not path.is_file()] + if missing: + raise FileNotFoundError( + "Cloth twist asset package is incomplete; missing: " + + ", ".join(str(path) for path in missing) + ) + + vertices: list[list[float]] = [] + texcoords: list[list[float]] = [] + faces: list[list[int]] = [] + face_uvs: list[list[int]] = [] + for line in source_mesh.read_text(encoding="utf-8").splitlines(): + if line.startswith("v "): + _, x, y, z = line.split()[:4] + vertices.append([float(x), float(y), float(z)]) + elif line.startswith("vt "): + _, u, v = line.split()[:3] + texcoords.append([float(u), float(v)]) + elif line.startswith("f "): + references = [item.split("/") for item in line.split()[1:]] + if len(references) != 3: + raise RuntimeError(f"{source_mesh} must contain triangle faces only.") + faces.append([int(reference[0]) - 1 for reference in references]) + face_uvs.append( + [ + int(reference[1]) - 1 if len(reference) > 1 and reference[1] else -1 + for reference in references + ] + ) + + vertices_array = np.asarray(vertices, dtype=np.float32) * MESH_SCALE + triangles = np.asarray(faces, dtype=np.int32) + expected_vertices = GRID_SIZE * GRID_SIZE + expected_faces = 2 * (GRID_SIZE - 1) * (GRID_SIZE - 1) + if vertices_array.shape != (expected_vertices, 3) or triangles.shape != ( + expected_faces, + 3, + ): + raise RuntimeError( + "cloth_twist_square_cloth.obj does not match the expected " + f"{GRID_SIZE} x {GRID_SIZE} topology." + ) + + vertex_uvs = np.full((len(vertices_array), 2), np.nan, dtype=np.float32) + for face, face_uv in zip(triangles, face_uvs, strict=True): + for vertex_index, texcoord_index in zip(face, face_uv, strict=True): + if texcoord_index < 0: + continue + uv = np.asarray(texcoords[texcoord_index], dtype=np.float32) + if np.isnan(vertex_uvs[vertex_index, 0]): + vertex_uvs[vertex_index] = uv + elif not np.allclose(vertex_uvs[vertex_index], uv, atol=1.0e-6): + raise RuntimeError( + f"{source_mesh} maps multiple UVs to vertex {vertex_index}." + ) + vertex_uvs[np.isnan(vertex_uvs[:, 0])] = 0.0 + return texture_path, vertices_array, triangles, vertex_uvs + + +def build_twist_trajectory( + vertices: np.ndarray, + frame_count: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Build fixed-node flags and opposite edge-rotation offsets. + + Args: + vertices: Scaled source vertices before the configured cloth pose. + frame_count: Number of outer 60 Hz simulation frames. + + Returns: + Shared node indices, per-node particle flags, and batched offsets. + """ + left_edge = np.asarray( + [GRID_SIZE - 1 + row * GRID_SIZE for row in range(GRID_SIZE)], + dtype=np.int32, + ) + right_edge = np.asarray( + [row * GRID_SIZE for row in range(GRID_SIZE)], + dtype=np.int32, + ) + node_indices = np.concatenate((left_edge, right_edge)) + + particle_flags = np.ones(len(vertices), dtype=np.int32) + particle_flags[node_indices] = 0 + + angle = np.pi / 2.0 + cloth_rotation = np.asarray( + [ + [np.cos(angle), -np.sin(angle), 0.0], + [np.sin(angle), np.cos(angle), 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + posed_vertices = vertices @ cloth_rotation.T + selected_positions = posed_vertices[node_indices] + rotation_axes = np.asarray( + [[0.0, 1.0, 0.0]] * len(left_edge) + [[0.0, -1.0, 0.0]] * len(right_edge), + dtype=np.float32, + ) + roots = ( + np.sum(selected_positions * rotation_axes, axis=1, keepdims=True) + * rotation_axes + ) + radial_vectors = selected_positions - roots + axis_cross_radial = np.cross(rotation_axes, radial_vectors) + axis_dot_radial = np.sum( + rotation_axes * radial_vectors, + axis=1, + keepdims=True, + ) + + sample_count = frame_count * NUM_SUBSTEPS + times = np.minimum( + np.arange(sample_count, dtype=np.float32) / (FPS * NUM_SUBSTEPS), + ROTATION_END_TIME, + ) + theta = times * ROTATION_ANGULAR_VELOCITY + cosine = np.cos(theta)[:, None, None] + sine = np.sin(theta)[:, None, None] + rotated_radial = ( + cosine * radial_vectors[None] + + sine * axis_cross_radial[None] + + (1.0 - cosine) * rotation_axes[None] * axis_dot_radial[None] + ) + target_positions = roots[None] + rotated_radial + offsets = target_positions - selected_positions[None] + return node_indices, particle_flags, offsets[None].astype(np.float32) + + +def initialize_simulation(args: argparse.Namespace) -> SimulationManager: + """Create the zero-gravity Newton VBD simulation manager.""" + cfg = SimulationManagerCfg( + width=1920, + height=1080, + headless=args.headless, + device=args.device, + gpu_id=args.gpu_id, + num_envs=args.num_envs, + arena_space=args.arena_space, + render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=NewtonPhysicsCfg( + physics_dt=1.0 / FPS, + device=args.device, + gravity=(0.0, 0.0, 0.0), + num_substeps=NUM_SUBSTEPS, + use_cuda_graph=False, + solver_cfg={ + "solver_type": "vbd", + "iterations": SOLVER_ITERATIONS, + "particle_enable_self_contact": True, + "particle_self_contact_radius": 0.002, + "particle_self_contact_margin": 0.0035, + "particle_enable_tile_solve": True, + "soft_contact_ke": 1.0e3, + "soft_contact_kd": 1.0e-1, + "soft_contact_mu": 0.2, + }, + ), + visualization=visualization_cfg_from_args(args), + ) + return SimulationManager(cfg) + + +def create_cloth( + sim: SimulationManager, + texture_path: Path, + vertices: np.ndarray, + triangles: np.ndarray, + uv_coords: np.ndarray, + particle_flags: np.ndarray, +) -> ClothObject: + """Declare the textured cloth with reference VBD material parameters.""" + return sim.add_cloth_object( + ClothObjectCfg( + uid="twist_cloth", + shape=MeshCfg( + vertices=vertices, + triangles=triangles, + uv_coords=uv_coords, + visual_material=VisualMaterialCfg( + uid="twist_cloth_material", + base_color=[1.0, 1.0, 1.0, 1.0], + base_color_texture=str(texture_path), + roughness=0.8, + metallic=0.0, + ), + ), + init_pos=CLOTH_POSITION, + init_rot=(0.0, 0.0, 90.0), + particle_flags=particle_flags, + physical_attr=ClothPhysicalAttributesCfg( + density=0.2, + tri_ke=1.0e3, + tri_ka=1.0e3, + tri_kd=2.0e-4, + edge_ke=1.0e-3, + edge_kd=1.0e-2, + ), + ) + ) + + +def configure_window_camera(sim: SimulationManager) -> None: + """Frame the vertical cloth in the native viewer.""" + window = sim.get_world().get_windows() + if window is not None: + window.set_look_at( + eye=np.asarray([2.25, 0.0, CLOTH_POSITION[2]], dtype=np.float32), + look_at=np.asarray(CLOTH_POSITION, dtype=np.float32), + up=np.asarray([0.0, 0.0, 1.0], dtype=np.float32), + ) + + +def main() -> None: + """Create the scene and execute the finite cloth-twist trajectory.""" + args = parse_arguments() + texture_path, vertices, triangles, uv_coords = prepare_cloth_asset( + args.cloth_material + ) + node_indices, particle_flags, offsets = build_twist_trajectory( + vertices, + args.iterations, + ) + sim = initialize_simulation(args) + + try: + cloth = create_cloth( + sim, + texture_path, + vertices, + triangles, + uv_coords, + particle_flags, + ) + sim.register_kinematic_nodal_trajectory( + cloth.uid, + node_indices, + offsets, + rebuild_self_contact_bvh=True, + ) + sim.prepare() + + if not args.headless and sim.open_window(): + configure_window_camera(sim) + + particle_count = cloth.get_default_nodal_state().shape[1] + logger.log_info( + f"Running cloth twist for {args.iterations} frames at {FPS} Hz " + f"with {particle_count} particles." + ) + for frame in range(args.iterations): + sim.update(step=1) + if frame % 50 == 0: + logger.log_info( + f"Frame {frame}/{args.iterations}, sim_time={frame / FPS:.2f}s" + ) + logger.log_info("Cloth twist simulation complete.") + except KeyboardInterrupt: + logger.log_info("\nExit") + finally: + sim.destroy() + + +if __name__ == "__main__": + main() diff --git a/examples/sim/demo/fold_tshirt.py b/examples/sim/demo/fold_tshirt.py new file mode 100644 index 000000000..d2c888bdf --- /dev/null +++ b/examples/sim/demo/fold_tshirt.py @@ -0,0 +1,702 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Fold a live Newton cloth T-shirt with a trajectory-driven DexForce W1.""" + +from __future__ import annotations + +import argparse +import math +import time +from pathlib import Path + +import numpy as np + +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 ( + ArticulationRootPropertiesCfg, + ClothObjectCfg, + ClothPhysicalAttributesCfg, + CollisionPropertiesCfg, + JointDrivePropertiesCfg, + JointDynamicsPropertiesCfg, + NewtonCollisionPropertiesCfg, + NewtonPhysicsCfg, + NewtonRigidBodyMaterialCfg, + NewtonRigidBodyPhysicsCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, + RobotCfg, + RenderCfg, +) +from embodichain.lab.sim.material import VisualMaterialCfg +from embodichain.lab.sim.objects import ClothObject, RigidObject, Robot +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg +from embodichain.lab.visualization import visualization_cfg_from_args +from embodichain.utils import logger + +DEFAULT_DT = 1.0 / 60.0 +DEFAULT_TRAJECTORY_TIME_SCALE = 4.0 +NUM_SUBSTEPS = 9 +SOLVER_ITERATIONS = 20 +FPS_LOG_INTERVAL = 120 + +ASSET_DATASET = "DeformableDemoData" +ANNIVERSARY_MATERIAL = "anniversary" +ANNIVERSARY_VISUAL_OBJ = "shirt_with_front_anniversary_decal_fold_atlas.obj" +ANNIVERSARY_TEXTURE = "shirt_with_front_anniversary_decal_fold_atlas.png" +TEXTURE_FILE_MAP = { + "mianbu": "mianbu.png", + "shabu": "shabu.png", + "mabu": "mabu.png", + "pige": "pige.png", + "jinduan": "jinduan.png", + "niuzai": "niuzai.png", + "wenli": "wenli.png", +} + +TABLE_POSITION = (0.55, 0.0, 1.15) +TABLE_SIZE = (0.52, 1.24, 0.05) +GROUND_POSITION = (0.0, 0.0, -0.01) +GROUND_SIZE = (8.0, 8.0, 0.02) +SHIRT_POSITION = (0.55, 0.0, 1.189) +SHIRT_SCALE = 0.0080 * 0.8 +CONTACT_SCRIPT_TRANSITIONS = np.asarray( + [4.21, 16.8, 19.0, 22.18, 27.4, 31.4], + dtype=np.float32, +) +CONTACT_SCRIPT_TO_SIMULATION_OFFSET = 1.2 + + +def parse_arguments() -> argparse.Namespace: + """Parse and validate command-line arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--steps", + type=int, + default=None, + help="Physics frames to run; defaults to the complete cached trajectory.", + ) + parser.add_argument( + "--dt", + type=float, + default=DEFAULT_DT, + help="Outer EmbodiChain physics timestep in seconds.", + ) + parser.add_argument( + "--trajectory-time-scale", + type=float, + default=DEFAULT_TRAJECTORY_TIME_SCALE, + help="Playback-speed multiplier for the cached W1 trajectory.", + ) + parser.add_argument( + "--static-w1", + action="store_true", + help="Keep W1 at the first trajectory pose for scene inspection.", + ) + parser.add_argument( + "--disable-w1-collision", + action="store_true", + help="Disable particle collision on all W1 links for debugging.", + ) + parser.add_argument( + "--disable-cuda-graph", + action="store_true", + help="Use direct Newton stepping and enable particle-friction scheduling.", + ) + parser.add_argument( + "--real-time", + action="store_true", + help="Sleep after each frame to approximate wall-clock playback.", + ) + parser.add_argument( + "--cloth-material", + choices=(ANNIVERSARY_MATERIAL, *TEXTURE_FILE_MAP), + default=ANNIVERSARY_MATERIAL, + help=( + "Use the authored anniversary atlas or one of the tiled fabric " + "textures packaged with the reference scene." + ), + ) + parser.add_argument( + "--w1-urdf", + type=Path, + default=None, + help="Optional W1 URDF override.", + ) + parser.add_argument( + "--trajectory", + type=Path, + default=None, + help="Optional cached W1 trajectory override.", + ) + parser.set_defaults(device="cuda", physics="newton", renderer="rt") + args = parser.parse_args() + + if args.physics != "newton": + parser.error("T-shirt folding requires --physics newton.") + if not str(args.device).startswith("cuda"): + parser.error("Newton MJVBD cloth simulation requires a CUDA device.") + if args.num_envs != 1: + parser.error("This trajectory scene currently supports --num_envs 1.") + if not math.isfinite(args.dt) or args.dt <= 0.0: + parser.error("--dt must be finite and positive.") + if ( + not math.isfinite(args.trajectory_time_scale) + or args.trajectory_time_scale <= 0.0 + ): + parser.error("--trajectory-time-scale must be finite and positive.") + if args.steps is not None and args.steps < 0: + parser.error("--steps must be non-negative.") + return args + + +def resolve_assets( + args: argparse.Namespace, +) -> tuple[Path, Path, Path, Path, Path | None, Path | None]: + """Resolve the W1, trajectory, simulation mesh, and selected visual assets.""" + asset_root = Path(get_data_path(f"{ASSET_DATASET}/fold_tshirt")) + texture_root = asset_root.parent / "textures" + + urdf_path = asset_root / "W1-hand-obj" / "DexforceW1V021_visual_collision.urdf" + trajectory_path = asset_root / "fold_tshirt.npz" + shirt_asset_root = asset_root / "shirt_front_decal_asset" + shirt_mesh_path = shirt_asset_root / "shirt_mesh.txt" + ground_texture_path = texture_root / "ground.png" + visual_mesh_path: Path | None = None + texture_path: Path | None = None + + if args.w1_urdf is not None: + urdf_path = args.w1_urdf.expanduser().resolve() + if args.trajectory is not None: + trajectory_path = args.trajectory.expanduser().resolve() + if args.cloth_material == ANNIVERSARY_MATERIAL: + visual_mesh_path = shirt_asset_root / ANNIVERSARY_VISUAL_OBJ + texture_path = shirt_asset_root / ANNIVERSARY_TEXTURE + else: + texture_path = texture_root / TEXTURE_FILE_MAP[args.cloth_material] + + required = [ + urdf_path, + trajectory_path, + shirt_mesh_path, + ground_texture_path, + ] + if visual_mesh_path is not None: + required.extend( + [ + visual_mesh_path, + visual_mesh_path.with_suffix(".mtl"), + ] + ) + if texture_path is not None: + required.append(texture_path) + missing = [str(path) for path in required if not path.is_file()] + if missing: + raise FileNotFoundError( + "W1 fold asset package is incomplete; missing: " + ", ".join(missing) + ) + return ( + urdf_path, + trajectory_path, + shirt_mesh_path, + ground_texture_path, + visual_mesh_path, + texture_path, + ) + + +def load_shirt_mesh(path: Path) -> tuple[np.ndarray, np.ndarray]: + """Load and normalize the reference shirt simulation mesh.""" + vertices: list[list[float]] = [] + triangles: list[list[int]] = [] + section: str | None = None + for raw_line in path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line: + continue + if line == "shirt_vertices": + section = "vertices" + continue + if line == "shirt_indices": + section = "indices" + continue + if section is None or ":" not in line: + raise RuntimeError(f"Unexpected shirt mesh line: {line!r}.") + _, raw_values = line.split(":", 1) + values = raw_values.split() + if len(values) < 3: + raise RuntimeError(f"Incomplete shirt mesh line: {line!r}.") + if section == "vertices": + vertices.append([float(value) for value in values[:3]]) + else: + triangles.append([int(value) for value in values[:3]]) + + vertex_array = np.asarray(vertices, dtype=np.float32) + triangle_array = np.asarray(triangles, dtype=np.int32) + if vertex_array.ndim != 2 or vertex_array.shape[1:] != (3,) or not len(vertices): + raise RuntimeError(f"{path} contains no valid shirt vertices.") + if ( + triangle_array.ndim != 2 + or triangle_array.shape[1:] != (3,) + or not len(triangles) + ): + raise RuntimeError(f"{path} contains no valid shirt triangles.") + if np.any(triangle_array < 0) or np.any(triangle_array >= len(vertex_array)): + raise RuntimeError(f"{path} contains out-of-range triangle indices.") + + minimum = vertex_array.min(axis=0) + maximum = vertex_array.max(axis=0) + vertex_array[:, :2] -= 0.5 * (minimum[:2] + maximum[:2]) + vertex_array[:, 2] -= minimum[2] + return vertex_array * SHIRT_SCALE, triangle_array + + +def vertex_normals(vertices: np.ndarray, triangles: np.ndarray) -> np.ndarray: + """Compute area-weighted per-vertex normals.""" + normals = np.zeros_like(vertices) + for triangle in triangles: + normal = np.cross( + vertices[triangle[1]] - vertices[triangle[0]], + vertices[triangle[2]] - vertices[triangle[0]], + ) + normals[triangle] += normal + lengths = np.linalg.norm(normals, axis=1, keepdims=True) + valid = lengths[:, 0] > 0.0 + normals[valid] /= lengths[valid] + return normals + + +def planar_uv(vertices: np.ndarray) -> np.ndarray: + """Generate the tiled planar UVs used by the fabric material variants.""" + uv_coords = vertices[:, :2].copy() + uv_coords -= uv_coords.min(axis=0) + uv_coords /= np.maximum(uv_coords.max(axis=0), 1.0e-8) + uv_coords[:, 1] = 1.0 - uv_coords[:, 1] + return uv_coords * 2.0 + + +def load_trajectory(path: Path, time_scale: float) -> tuple[np.ndarray, float]: + """Load the cached W1 public-qpos trajectory and compute playback FPS.""" + with np.load(path) as archive: + trajectory = np.asarray(archive["robot_qpos"], dtype=np.float32) + source_dt = float(np.asarray(archive["dt"]).reshape(-1)[0]) + if trajectory.ndim != 2 or trajectory.shape[0] == 0 or trajectory.shape[1] < 3: + raise ValueError( + f"Expected a non-empty trajectory with shape [frames, dof], got " + f"{trajectory.shape}." + ) + if not np.isfinite(trajectory).all(): + raise ValueError("W1 trajectory must contain only finite values.") + if not math.isfinite(source_dt) or source_dt <= 0.0: + raise ValueError("W1 trajectory dt must be finite and positive.") + + # The cache already follows the public Spawn qpos order. Keep all columns + # intact while locking the three leg joints so the torso stays at table height. + trajectory = np.ascontiguousarray(trajectory, dtype=np.float32).copy() + trajectory[:, :3] = 0.0 + trajectory_fps = (1.0 / source_dt) * time_scale / DEFAULT_TRAJECTORY_TIME_SCALE + return trajectory, trajectory_fps + + +def initialize_simulation( + args: argparse.Namespace, + *, + use_cuda_graph: bool, +) -> SimulationManager: + """Create the EmbodiChain manager with the reference MJVBD settings.""" + cfg = SimulationManagerCfg( + width=1920, + height=1080, + headless=args.headless, + device=args.device, + gpu_id=args.gpu_id, + num_envs=args.num_envs, + arena_space=args.arena_space, + render_cfg=RenderCfg(renderer=args.renderer, spp=1), + physics_cfg=NewtonPhysicsCfg( + physics_dt=args.dt, + device=args.device, + num_substeps=NUM_SUBSTEPS, + use_cuda_graph=use_cuda_graph, + solver_cfg={ + "solver_type": "mjvbd", + "iterations": SOLVER_ITERATIONS, + "particle_enable_self_contact": True, + "particle_self_contact_radius": 0.002, + "particle_self_contact_margin": 0.002, + "particle_topological_contact_filter_threshold": 1, + "particle_rest_shape_contact_exclusion_radius": 0.005, + "particle_vertex_contact_buffer_size": 96, + "particle_edge_contact_buffer_size": 128, + "particle_collision_detection_interval": -1, + "self_contact_bvh_rebuild_interval_frames": 15, + "rigid_contact_max": 0, + "step_rigid_bodies": False, + "soft_contact_margin": 0.008, + "soft_contact_ke": 3.0e5, + "soft_contact_kd": 5.0e-2, + "soft_contact_mu": 0.5, + }, + ), + visualization=visualization_cfg_from_args(args), + ) + return SimulationManager(cfg) + + +def create_table(sim: SimulationManager) -> RigidObject: + """Declare the static folding table.""" + return sim.add_rigid_object( + RigidObjectCfg( + uid="table", + shape=CubeCfg( + size=list(TABLE_SIZE), + visual_material=VisualMaterialCfg( + uid="fold_table_material", + base_color=[0.35, 0.42, 0.48, 1.0], + roughness=0.7, + metallic=0.0, + ), + ), + attrs=RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg(collision_enabled=True), + material_props=RigidBodyMaterialCfg( + static_friction=0.5, + dynamic_friction=0.5, + ), + newton_props=NewtonRigidBodyPhysicsCfg( + collision_props=NewtonCollisionPropertiesCfg( + has_particle_collision=True, + ), + material_props=NewtonRigidBodyMaterialCfg( + ke=5.0e5, + kd=1.0e-6, + ), + ), + ), + body_type="static", + init_pos=TABLE_POSITION, + ) + ) + + +def create_ground(sim: SimulationManager, texture_path: Path) -> RigidObject: + """Declare the textured static ground used by the reference scene.""" + return sim.add_rigid_object( + RigidObjectCfg( + uid="ground", + shape=CubeCfg( + size=list(GROUND_SIZE), + visual_material=VisualMaterialCfg( + uid="fold_ground_material", + base_color=[1.0, 1.0, 1.0, 1.0], + base_color_texture=str(texture_path), + roughness=0.65, + metallic=0.0, + ), + ), + attrs=RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg(collision_enabled=True), + material_props=RigidBodyMaterialCfg( + static_friction=0.5, + dynamic_friction=0.5, + ), + newton_props=NewtonRigidBodyPhysicsCfg( + collision_props=NewtonCollisionPropertiesCfg( + has_particle_collision=True, + ), + material_props=NewtonRigidBodyMaterialCfg( + ke=5.0e5, + kd=1.0e-6, + ), + ), + ), + body_type="static", + init_pos=GROUND_POSITION, + ) + ) + + +def create_w1( + sim: SimulationManager, + urdf_path: Path, + initial_qpos: np.ndarray, + *, + particle_collision_enabled: bool, +) -> Robot: + """Declare the fixed-base W1 using its exact source URDF.""" + robot = sim.add_robot( + RobotCfg( + uid="w1", + fpath=str(urdf_path), + asset_physics_mode="overlay", + articulation_props=ArticulationRootPropertiesCfg( + fixed_base=True, + self_collision_enabled=False, + ), + # Preserve source drive modes; the runtime control writes kinematic + # joint state directly at every Newton substep. + drive_pros=JointDrivePropertiesCfg(), + joint_props=JointDynamicsPropertiesCfg( + # Several hand joints author zero limits in the URDF, which + # produce invalid MuJoCo actfrcrange values without this overlay. + max_effort=180.0, + max_velocity=4.0, + ), + attrs=RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg(collision_enabled=True), + material_props=RigidBodyMaterialCfg( + static_friction=0.25, + dynamic_friction=0.25, + ), + newton_props=NewtonRigidBodyPhysicsCfg( + collision_props=NewtonCollisionPropertiesCfg( + has_particle_collision=particle_collision_enabled, + ), + material_props=NewtonRigidBodyMaterialCfg( + ke=3.0e5, + kd=1.0e-4, + ), + ), + ), + init_qpos=initial_qpos, + build_pk_chain=False, + ) + ) + if robot is None: + raise RuntimeError("Failed to declare the DexForce W1 robot.") + return robot + + +def create_shirt( + sim: SimulationManager, + vertices: np.ndarray, + triangles: np.ndarray, + *, + visual_mesh_path: Path | None, + texture_path: Path | None, +) -> ClothObject: + """Declare the low-resolution cloth and its independently bound visual mesh.""" + if visual_mesh_path is not None: + # The OBJ carries seam-duplicated vertices, authored UVs, and its own + # double-sided MTL. Leaving visual_material unset preserves that MTL. + visual_shape = MeshCfg(fpath=str(visual_mesh_path)) + else: + if texture_path is None: + raise ValueError("A fabric texture is required without an atlas mesh.") + visual_shape = MeshCfg( + vertices=vertices, + triangles=triangles, + normals=vertex_normals(vertices, triangles), + uv_coords=planar_uv(vertices), + visual_material=VisualMaterialCfg( + uid="fold_shirt_material", + base_color=[1.0, 1.0, 1.0, 1.0], + base_color_texture=str(texture_path), + roughness=0.8, + metallic=0.0, + ), + ) + + return sim.add_cloth_object( + ClothObjectCfg( + uid="shirt", + shape=MeshCfg(vertices=vertices, triangles=triangles), + visual_shape=visual_shape, + visual_binding_mode="nearest_vertex", + init_pos=SHIRT_POSITION, + init_rot=(0.0, 0.0, -90.0), + particle_radius=0.008, + physical_attr=ClothPhysicalAttributesCfg( + density=200.0, + tri_ke=1.5e3, + tri_ka=1.5e3, + tri_kd=1.0e-5, + edge_ke=1.2, + edge_kd=0.1, + ), + ) + ) + + +def register_runtime_controls( + sim: SimulationManager, + trajectory: np.ndarray, + trajectory_fps: float, + trajectory_time_scale: float, + *, + static_w1: bool, + use_cuda_graph: bool, +) -> None: + """Register the folding trajectory and phase-dependent contact materials.""" + transition_times = ( + CONTACT_SCRIPT_TRANSITIONS + CONTACT_SCRIPT_TO_SIMULATION_OFFSET + ) / trajectory_time_scale + contact_times = np.concatenate([np.zeros(1, dtype=np.float32), transition_times]) + + def friction_track(values: tuple[float, ...]) -> tuple[tuple[float, float], ...]: + return tuple( + (float(sample_time), value) + for sample_time, value in zip(contact_times, values, strict=True) + ) + + particle_friction = friction_track((0.5, 1.2, 0.0, 0.5, 1.2, 0.0, 0.5)) + w1_friction = friction_track((0.25, 1.2, 0.0, 0.25, 1.2, 0.0, 0.25)) + table_friction = friction_track((0.5, 0.12, 0.5, 0.5, 0.12, 0.5, 0.5)) + + # This global particle control is host-side. Graph mode keeps the initial + # soft_contact_mu while the graph-compatible rigid schedules still run. + if not use_cuda_graph: + sim.register_particle_contact_material_schedule( + {"dynamic_friction": particle_friction} + ) + sim.register_contact_material_schedule( + "w1", + {"dynamic_friction": w1_friction}, + ) + sim.register_contact_material_schedule( + "table", + {"dynamic_friction": table_friction}, + ) + sim.register_contact_material_schedule( + "ground", + {"dynamic_friction": table_friction}, + ) + if not static_w1: + sim.register_kinematic_joint_trajectory( + "w1", + trajectory[None], + fps=trajectory_fps, + ) + + +def configure_window_camera(sim: SimulationManager) -> None: + """Frame both W1 arms and the shirt on the table.""" + window = sim.get_world().get_windows() + if window is not None: + window.set_look_at( + eye=np.asarray([1.15, -2.10, 1.65], dtype=np.float32), + look_at=np.asarray([0.55, 0.0, 1.18], dtype=np.float32), + up=np.asarray([0.0, 0.0, 1.0], dtype=np.float32), + ) + + +def main() -> None: + """Build the EmbodiChain scene and play the complete folding trajectory.""" + args = parse_arguments() + ( + urdf_path, + trajectory_path, + shirt_mesh_path, + ground_texture_path, + visual_mesh_path, + texture_path, + ) = resolve_assets(args) + vertices, triangles = load_shirt_mesh(shirt_mesh_path) + trajectory, trajectory_fps = load_trajectory( + trajectory_path, + args.trajectory_time_scale, + ) + if args.steps is None: + args.steps = int(math.ceil(len(trajectory) / trajectory_fps / args.dt)) + + use_cuda_graph = not args.disable_cuda_graph + sim = initialize_simulation(args, use_cuda_graph=use_cuda_graph) + try: + # Replace EmbodiChain's default grid plane with the reference wood floor. + sim.set_ground_plane_visibility(False) + create_ground(sim, ground_texture_path) + create_table(sim) + robot = create_w1( + sim, + urdf_path, + trajectory[0], + particle_collision_enabled=not args.disable_w1_collision, + ) + shirt = create_shirt( + sim, + vertices, + triangles, + visual_mesh_path=visual_mesh_path, + texture_path=texture_path, + ) + register_runtime_controls( + sim, + trajectory, + trajectory_fps, + args.trajectory_time_scale, + static_w1=args.static_w1, + use_cuda_graph=use_cuda_graph, + ) + sim.prepare() + + if trajectory.shape[1] != robot.dof: + raise ValueError( + f"W1 trajectory has {trajectory.shape[1]} columns, but the " + f"articulation has {robot.dof} DOFs." + ) + if not args.headless and sim.open_window(): + sim.set_emission_light([1.0, 1.0, 1.0], 90.0) + configure_window_camera(sim) + + particle_count = shirt.get_default_nodal_state().shape[1] + logger.log_info( + "Running W1 T-shirt fold | " + f"driver={'static' if args.static_w1 else 'kinematic-trajectory'} | " + f"cuda_graph={use_cuda_graph} | " + f"frames={len(trajectory)} | trajectory_fps={trajectory_fps:.3f} | " + f"dof={robot.dof} | cloth_particles={particle_count} | " + f"material={args.cloth_material}" + ) + + fps_window_start: float | None = None + fps_window_steps = 0 + for frame in range(args.steps): + frame_start = time.perf_counter() + sim.update(step=1) + if args.real_time: + elapsed = time.perf_counter() - frame_start + time.sleep(max(0.0, args.dt - elapsed)) + + frame_end = time.perf_counter() + if fps_window_start is None: + # Exclude one-time Warp compilation and CUDA Graph capture. + fps_window_start = frame_end + else: + fps_window_steps += 1 + if (frame + 1) % FPS_LOG_INTERVAL == 0 or frame == args.steps - 1: + elapsed = frame_end - fps_window_start + fps = ( + fps_window_steps / elapsed + if elapsed > 0.0 and fps_window_steps > 0 + else 0.0 + ) + logger.log_info(f"Frame {frame + 1}/{args.steps}, FPS={fps:.1f}") + fps_window_start = frame_end + fps_window_steps = 0 + logger.log_info("W1 T-shirt folding simulation complete.") + except KeyboardInterrupt: + logger.log_info("\nExit") + finally: + sim.destroy() + + +if __name__ == "__main__": + main() diff --git a/examples/sim/demo/pick_up_cloth.py b/examples/sim/demo/pick_up_cloth.py index f4e60a8c8..3cfec4036 100644 --- a/examples/sim/demo/pick_up_cloth.py +++ b/examples/sim/demo/pick_up_cloth.py @@ -14,45 +14,54 @@ # limitations under the License. # ---------------------------------------------------------------------------- -""" -This script demonstrates the creation and simulation of a robot with a soft object, -and performs a pressing task in a simulated environment. -""" +"""Pick up a cloth with a UR10 gripper using the Newton MJVBD solver.""" from __future__ import annotations import argparse +import os +import tempfile +from collections.abc import Sequence + import numpy as np -import time import open3d as o3d import torch +from scipy.spatial.transform import Rotation -from dexsim.utility.path import get_resources_data_path - -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.objects import Robot, SoftObject -from embodichain.lab.sim.utility.action_utils import interpolate_with_distance -from embodichain.lab.sim.shapes import MeshCfg 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 +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( - RenderCfg, - physics_cfg_for_backend, - RigidObjectCfg, - RigidBodyAttributesCfg, - LightCfg, ClothObjectCfg, ClothPhysicalAttributesCfg, + MassPropertiesCfg, + NewtonCollisionPropertiesCfg, + NewtonPhysicsCfg, + NewtonRigidBodyMaterialCfg, + NewtonRigidBodyPhysicsCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, + RenderCfg, ) +from embodichain.lab.sim.objects import ClothObject, RigidObject, Robot from embodichain.lab.sim.robots import URRobotCfg -import os -from embodichain.lab.sim.shapes import MeshCfg, CubeCfg -import tempfile -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg +from embodichain.lab.sim.utility.action_utils import interpolate_with_nums +from embodichain.lab.visualization import visualization_cfg_from_args +CLOTH_SIZE = 0.3 +CLOTH_GRID_CELLS = 50 +CLOTH_PARTICLE_RADIUS = 0.003 +CLOTH_RIGID_CONTACT_KE = 1.0e6 +CLOTH_RIGID_CONTACT_KD = 5.0e-2 +CLOTH_RIGID_CONTACT_MU = 0.5 -def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]): + +def create_robot( + sim: SimulationManager, + position: Sequence[float] = (0.0, 0.0, 0.0), +) -> Robot: """ Create and configure a robot with an arm and a dexterous hand in the simulation. @@ -77,6 +86,32 @@ def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]): "damping": {"FINGER[1-2]": 1e1}, "max_effort": {"FINGER[1-2]": 1e3}, "drive_type": "force", + "target_mode": "position_velocity", + }, + "link_attrs": { + "gripper_collision": { + "link_names_expr": ["hand_base_link", "finger[1-2]"], + "attrs": { + "collision_props": { + "collision_enabled": True, + # Detect the thin cloth before it reaches the mesh. + "contact_offset": 0.008, + "rest_offset": 0.002, + }, + "material_props": { + "dynamic_friction": 2.0, + }, + "newton_props": { + "collision_props": { + "has_particle_collision": True, + }, + "material_props": { + "ke": 3.0e5, + "kd": 1.0e-4, + }, + }, + }, + }, }, "control_parts": { "hand": ["FINGER[1-2]"], @@ -107,29 +142,39 @@ def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]): return sim.add_robot(cfg=cfg) -def create_padding_box(sim: SimulationManager): +def create_padding_box(sim: SimulationManager) -> RigidObject: padding_box_cfg = RigidObjectCfg( uid="padding_box", 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( + mass_props=MassPropertiesCfg(mass=1.0), + material_props=RigidBodyMaterialCfg( + static_friction=CLOTH_RIGID_CONTACT_MU, + dynamic_friction=CLOTH_RIGID_CONTACT_MU, + restitution=0.01, + ), + newton_props=NewtonRigidBodyPhysicsCfg( + collision_props=NewtonCollisionPropertiesCfg( + has_particle_collision=True, + ), + material_props=NewtonRigidBodyMaterialCfg( + ke=CLOTH_RIGID_CONTACT_KE, + kd=CLOTH_RIGID_CONTACT_KD, + ), + ), ), body_type="kinematic", init_pos=[0.5, 0.0, 0.026], init_rot=[0.0, 0.0, 0.0], ) - padding_box = sim.add_rigid_object(cfg=padding_box_cfg) - return padding_box + return sim.add_rigid_object(cfg=padding_box_cfg) -def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): +def create_2d_grid_mesh( + width: float, height: float, nx: int = 1, ny: int = 1 +) -> tuple[torch.Tensor, torch.Tensor]: """Create a flat rectangle in the XY plane centered at `origin`. The rectangle is subdivided into an `nx` by `ny` grid (cells) and @@ -145,7 +190,7 @@ def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): # Vectorized vertex positions using PyTorch x_lin = torch.linspace(-w / 2.0, w / 2.0, steps=nx + 1, dtype=torch.float64) y_lin = torch.linspace(-h / 2.0, h / 2.0, steps=ny + 1, dtype=torch.float64) - yy, xx = torch.meshgrid(y_lin, x_lin) # shapes: (ny+1, nx+1) + yy, xx = torch.meshgrid(y_lin, x_lin, indexing="ij") xx_flat = xx.reshape(-1) yy_flat = yy.reshape(-1) zz_flat = torch.full_like(xx_flat, 0, dtype=torch.float64) @@ -163,8 +208,13 @@ def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): return verts, faces -def create_cloth(sim: SimulationManager): - cloth_verts, cloth_faces = create_2d_grid_mesh(width=0.3, height=0.3, nx=12, ny=12) +def create_cloth(sim: SimulationManager) -> ClothObject: + cloth_verts, cloth_faces = create_2d_grid_mesh( + width=CLOTH_SIZE, + height=CLOTH_SIZE, + nx=CLOTH_GRID_CELLS, + ny=CLOTH_GRID_CELLS, + ) cloth_mesh = o3d.geometry.TriangleMesh( vertices=o3d.utility.Vector3dVector(cloth_verts.to("cpu").numpy()), triangles=o3d.utility.Vector3iVector(cloth_faces.to("cpu").numpy()), @@ -178,36 +228,104 @@ def create_cloth(sim: SimulationManager): shape=MeshCfg(fpath=cloth_save_path), init_pos=[0.5, 0.0, 0.3], init_rot=[0, 0, 0], + # Keep the collision shell close to the rendered surface. A large + # radius can make the gripper carry cloth while visibly separated. + particle_radius=CLOTH_PARTICLE_RADIUS, physical_attr=ClothPhysicalAttributesCfg( - mass=0.01, - youngs=1e10, - poissons=0.4, - thickness=0.06, - bending_stiffness=0.01, - bending_damping=0.1, - dynamic_friction=0.95, - min_position_iters=30, + # Give gravity enough authority to produce fabric-like drape. + # Stretch remains firmer than bending so the cloth folds + # instead of behaving like an elastic sheet. + density=0.05, + tri_ke=2.0e2, + tri_ka=2.0e2, + tri_kd=1.0e-5, + edge_ke=0.005, + edge_kd=0.01, ), ) ) return cloth -def get_grasp_traj(sim: SimulationManager, robot: Robot, grasp_xpos: torch.Tensor): +def get_grasp_traj( + sim: SimulationManager, + robot: Robot, + grasp_xpos: torch.Tensor, +) -> torch.Tensor: + """Build the full robot trajectory without materializing the scene. + + MJVBD cloth manipulation uses an external kinematic articulation. The + trajectory therefore has to be registered before ``sim.prepare()`` so the + Newton runtime can interpolate joint poses and velocities at every + substep. + """ num_envs = sim.num_envs - rest_arm_qpos = robot.get_qpos("arm") + arm_joint_names = robot.cfg.control_parts["arm"] + arm_dof = len(arm_joint_names) + initial_qpos = torch.as_tensor( + robot.cfg.init_qpos, + dtype=torch.float32, + device=sim.device, + ).reshape(1, -1) + hand_dof = initial_qpos.shape[1] - arm_dof + if hand_dof <= 0: + raise ValueError("The robot trajectory requires at least one hand DOF.") + rest_arm_qpos = initial_qpos[:, :arm_dof].repeat(num_envs, 1) + + solver_cfg = robot.cfg.solver_cfg["arm"] + solver_cfg.joint_names = list(arm_joint_names) + if solver_cfg.urdf_path is None: + solver_cfg.urdf_path = robot.cfg.fpath + solver = solver_cfg.init_solver(device=sim.device) + + root_pose_value = robot.cfg.init_local_pose + if root_pose_value is None: + root_pose_value = np.eye(4, dtype=np.float32) + root_pose_value[:3, :3] = Rotation.from_euler( + "xyz", robot.cfg.init_rot, degrees=True + ).as_matrix() + root_pose_value[:3, 3] = np.asarray(robot.cfg.init_pos, dtype=np.float32) + root_pose = torch.as_tensor( + root_pose_value, + dtype=torch.float32, + device=sim.device, + ).reshape(1, 4, 4) + root_pose = root_pose.repeat(num_envs, 1, 1) + root_pose_inv = torch.linalg.inv(root_pose) approach_xpos = grasp_xpos.clone() - approach_xpos[:, 2, 3] += 0.04 - _, qpos_approach = robot.compute_ik( - pose=approach_xpos, joint_seed=rest_arm_qpos, name="arm" + approach_xpos[:, 2, 3] += 0.06 + approach_xpos = torch.bmm(root_pose_inv, approach_xpos) + local_grasp_xpos = torch.bmm(root_pose_inv, grasp_xpos) + approach_success, qpos_approach = solver.get_ik( + target_xpos=approach_xpos, + qpos_seed=rest_arm_qpos, + ) + grasp_success, qpos_grasp = solver.get_ik( + target_xpos=local_grasp_xpos, + qpos_seed=qpos_approach, ) - _, qpos_grasp = robot.compute_ik( - pose=grasp_xpos, joint_seed=qpos_approach, name="arm" + if not bool(torch.all(approach_success & grasp_success)): + failed_envs = torch.nonzero( + ~(approach_success & grasp_success), as_tuple=False + ).flatten() + raise RuntimeError(f"IK failed for environment indices {failed_envs.tolist()}.") + + hand_open_qpos = initial_qpos[:, arm_dof : arm_dof + hand_dof].repeat(num_envs, 1) + # First close around the cloth ridge while the padding box supports it. + # After lifting clear of the box, close once more so the fingers—not the + # support reaction—provide the sustained normal force for the grasp. + hand_pregrasp_qpos = torch.full( + (num_envs, hand_dof), + 0.012, + dtype=torch.float32, + device=sim.device, ) - hand_open_qpos = torch.tensor([0.00, 0.00], dtype=torch.float32, device=sim.device) - hand_close_qpos = torch.tensor( - [0.025, 0.025], dtype=torch.float32, device=sim.device + hand_grasp_qpos = torch.full( + (num_envs, hand_dof), + 0.024, + dtype=torch.float32, + device=sim.device, ) arm_trajectory = torch.cat( @@ -217,87 +335,159 @@ def get_grasp_traj(sim: SimulationManager, robot: Robot, grasp_xpos: torch.Tenso qpos_grasp[:, None, :], qpos_grasp[:, None, :], qpos_approach[:, None, :], + qpos_approach[:, None, :], rest_arm_qpos[:, None, :], ], dim=1, ) hand_trajectory = torch.cat( [ - hand_open_qpos[None, None, :].repeat(num_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(num_envs, 1, 1), - hand_open_qpos[None, None, :].repeat(num_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(num_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(num_envs, 1, 1), - hand_close_qpos[None, None, :].repeat(num_envs, 1, 1), + hand_open_qpos[:, None, :], + hand_open_qpos[:, None, :], + hand_open_qpos[:, None, :], + hand_pregrasp_qpos[:, None, :], + hand_pregrasp_qpos[:, None, :], + hand_grasp_qpos[:, None, :], + hand_grasp_qpos[:, None, :], ], dim=1, ) all_trajectory = torch.cat([arm_trajectory, hand_trajectory], dim=-1) - interp_trajectory = interpolate_with_distance( - trajectory=all_trajectory, interp_num=220, device=sim.device + # Keep the original wall-clock timing while supplying one waypoint per + # 100 Hz frame. Newton further interpolates each frame over 12 substeps. + interp_trajectory = interpolate_with_nums( + trajectory=all_trajectory, + interp_nums=torch.tensor([180, 90, 180, 90, 90, 180]), + device=sim.device, ) return interp_trajectory -def main(): +def register_kinematic_trajectory( + sim: SimulationManager, + trajectory: torch.Tensor, + settle_steps: int, +) -> None: + """Register the grasp trajectory after an initial settling hold. + + Args: + sim: Simulation manager that owns the target robot. + trajectory: Batched robot joint positions with shape + ``(num_envs, frames, dof)``. + settle_steps: Number of initial physics frames held at row zero. + """ + if settle_steps < 0: + raise ValueError("settle_steps must be non-negative.") + + # Row zero is the initial state and each simulation frame advances to the + # next row. Keep settle_steps + 1 identical rows so settling consumes no + # part of the grasp motion. + hold = trajectory[:, :1, :].repeat(1, settle_steps + 1, 1) + playback = torch.cat([hold, trajectory[:, 1:, :]], dim=1) + + # SimulationManager owns Spawn path expansion and the pre-prepare runtime + # control lifecycle. DexSim supplies the substep q/qdot interpolation and FK. + sim.register_kinematic_joint_trajectory("UR10", playback) + + +def main() -> None: """ Main function to demonstrate robot simulation. - This function initializes the simulation, creates the robot and other objects, - and performs the press softbody task. + This function initializes the simulation, creates the robot and cloth, + and executes the pick-up trajectory. """ parser = argparse.ArgumentParser( description="Create a simulation scene with SimulationManager" ) add_env_launcher_args_to_parser(parser) + parser.set_defaults(device="cuda", physics="newton") args = parser.parse_args() + if args.physics != "newton": + parser.error("Cloth requires --physics newton.") + if not str(args.device).startswith("cuda"): + parser.error("Newton MJVBD cloth simulation requires a CUDA device.") # Configure the simulation sim_cfg = SimulationManagerCfg( width=1920, height=1080, num_envs=args.num_envs, - headless=True, + arena_space=args.arena_space, + gpu_id=args.gpu_id, + headless=args.headless, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - device="cuda", + device=args.device, render_cfg=RenderCfg( renderer=args.renderer ), # Enable ray tracing for better visuals - physics_cfg=physics_cfg_for_backend(args.physics), + physics_cfg=NewtonPhysicsCfg( + num_substeps=12, + solver_cfg={ + "solver_type": "mjvbd", + "iterations": 24, + "particle_enable_self_contact": True, + "particle_self_contact_radius": 0.002, + "particle_self_contact_margin": 0.002, + "particle_topological_contact_filter_threshold": 1, + "particle_rest_shape_contact_exclusion_radius": 0.005, + "particle_vertex_contact_buffer_size": 96, + "particle_edge_contact_buffer_size": 128, + "particle_collision_detection_interval": -1, + "particle_enable_tile_solve": True, + "soft_contact_margin": 0.008, + "soft_contact_ke": CLOTH_RIGID_CONTACT_KE, + "soft_contact_kd": CLOTH_RIGID_CONTACT_KD, + "soft_contact_mu": CLOTH_RIGID_CONTACT_MU, + # Use the mixed material stiffness immediately instead of + # ramping new contacts from the low MJVBD default. + "rigid_contact_k_start": CLOTH_RIGID_CONTACT_KE, + "rigid_body_particle_contact_buffer_size": 512, + "rigid_contact_max": 0, + # The registered runtime control advances and interpolates the + # robot kinematically at every Newton substep. + "step_rigid_bodies": False, + "self_contact_bvh_rebuild_interval_frames": 1, + }, + ), visualization=visualization_cfg_from_args(args), ) # Create the simulation instance sim = SimulationManager(sim_cfg) - robot = create_robot(sim) - cloth = create_cloth(sim) - padding_box = create_padding_box(sim) - sim.prepare() - if not args.headless: - sim.open_window() - sim.update(step=10) # Let the cloth settle before interaction + try: + robot = create_robot(sim) + create_cloth(sim) + create_padding_box(sim) - grasp_xpos = torch.tensor( - [ + grasp_xpos = torch.tensor( [ - [-1, 0, 0, 0.5], - [0, 1, 0, 0], - [0, 0, -1, 0.075], - [0, 0, 0, 1], + [ + [-1, 0, 0, 0.5], + [0, 1, 0, 0], + [0, 0, -1, 0.075], + [0, 0, 0, 1], + ], ], - ], - dtype=torch.float32, - device=sim.device, - ) - grasp_xpos = grasp_xpos.repeat(sim.num_envs, 1, 1) - grab_traj = get_grasp_traj(sim, robot, grasp_xpos) - input("Press Enter to start grabing cloth...") - - n_waypoint = grab_traj.shape[1] - for i in range(n_waypoint): - robot.set_qpos(grab_traj[:, i, :]) - sim.update(step=3) - input("Press Enter to exit the simulation...") + dtype=torch.float32, + device=sim.device, + ) + grasp_xpos = grasp_xpos.repeat(sim.num_envs, 1, 1) + grab_traj = get_grasp_traj(sim, robot, grasp_xpos) + settle_steps = 100 + register_kinematic_trajectory(sim, grab_traj, settle_steps) + + sim.prepare() + if not args.headless: + sim.open_window() + sim.update(step=settle_steps) + input("Press Enter to start grabbing the cloth...") + + # The initial trajectory row is already active after settling. + sim.update(step=grab_traj.shape[1] - 1) + input("Press Enter to exit the simulation...") + finally: + sim.destroy() if __name__ == "__main__": diff --git a/examples/sim/demo/press_softbody.py b/examples/sim/demo/press_softbody.py index 017235276..2ec860203 100644 --- a/examples/sim/demo/press_softbody.py +++ b/examples/sim/demo/press_softbody.py @@ -14,93 +14,137 @@ # limitations under the License. # ---------------------------------------------------------------------------- -""" -This script demonstrates the creation and simulation of a robot with a soft object, -and performs a pressing task in a simulated environment. -""" +"""Press a soft cow with a UR10 using the Newton MJVBD solver.""" from __future__ import annotations import argparse + import numpy as np -import time import torch - from dexsim.utility.path import get_resources_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.visualization import visualization_cfg_from_args -from embodichain.lab.sim.objects import Robot, SoftObject -from embodichain.lab.sim.utility.action_utils import interpolate_with_distance -from embodichain.lab.sim.shapes import MeshCfg -from embodichain.data import get_data_path -from embodichain.utils import logger from embodichain.lab.sim.cfg import ( + NewtonCollisionPipelineCfg, + NewtonPhysicsCfg, RenderCfg, - physics_cfg_for_backend, - LightCfg, SoftObjectCfg, - SoftbodyVoxelAttributesCfg, SoftbodyPhysicalAttributesCfg, + SoftbodyVoxelAttributesCfg, ) -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.objects import Robot, SoftObject from embodichain.lab.sim.robots import URRobotCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.utility.action_utils import interpolate_with_nums +from embodichain.lab.visualization import visualization_cfg_from_args +from embodichain.utils import logger +PHYSICS_DT = 1.0 / 100.0 +NUM_SUBSTEPS = 6 +SOLVER_ITERATIONS = 8 +SETTLE_STEPS = 100 -def parse_arguments(): - """ - Parse command-line arguments to configure the simulation. +SOFT_CONTACT_MARGIN = 0.003 +SOFT_CONTACT_KE = 5.0e4 +SOFT_CONTACT_KD = 1.0e-3 +SOFT_CONTACT_MU = 1.0 - Returns: - argparse.Namespace: Parsed arguments including number of environments, device, and rendering options. - """ - parser = argparse.ArgumentParser( - description="Create and simulate a robot in SimulationManager" - ) - add_env_launcher_args_to_parser(parser) - return parser.parse_args() +COW_POSITION = (0.45, -0.1, 0.12) +PRESS_POSITION = (0.5, -0.1, 0.04) +APPROACH_HEIGHT = 0.015 -def initialize_simulation(args): - """ - Initialize the simulation environment based on the provided arguments. +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the demo.""" + parser = argparse.ArgumentParser(description=__doc__) + add_env_launcher_args_to_parser(parser) + parser.set_defaults(device="cuda", physics="newton") + args = parser.parse_args() + if args.physics != "newton": + parser.error("Soft bodies require --physics newton.") + if not str(args.device).startswith("cuda"): + parser.error("EmbodiChain soft bodies currently require a CUDA device.") + return args - Args: - args (argparse.Namespace): Parsed command-line arguments. - Returns: - SimulationManager: Configured simulation manager instance. - """ +def initialize_simulation(args: argparse.Namespace) -> SimulationManager: + """Create the Newton MJVBD simulation manager.""" config = SimulationManagerCfg( - headless=True, - device="cuda", - render_cfg=RenderCfg(renderer=args.renderer), - physics_cfg=physics_cfg_for_backend(args.physics), - physics_dt=1.0 / 100.0, + width=1920, + height=1080, + headless=args.headless, + device=args.device, + gpu_id=args.gpu_id, num_envs=args.num_envs, + arena_space=args.arena_space, + render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=NewtonPhysicsCfg( + physics_dt=PHYSICS_DT, + device=args.device, + num_substeps=NUM_SUBSTEPS, + solver_cfg={ + "solver_type": "mjvbd", + "iterations": SOLVER_ITERATIONS, + "particle_enable_self_contact": False, + "particle_self_contact_radius": 0.005, + "particle_self_contact_margin": 0.005, + "particle_topological_contact_filter_threshold": 3, + "particle_enable_tile_solve": True, + "rigid_body_particle_contact_buffer_size": 512, + "rigid_contact_k_start": SOFT_CONTACT_KE, + "rigid_contact_max": 0, + "soft_contact_margin": SOFT_CONTACT_MARGIN, + "soft_contact_ke": SOFT_CONTACT_KE, + "soft_contact_kd": SOFT_CONTACT_KD, + "soft_contact_mu": SOFT_CONTACT_MU, + # The registered trajectory updates the robot kinematically at + # every Newton substep; MJVBD only needs to solve the soft body. + "step_rigid_bodies": False, + }, + collision_cfg=NewtonCollisionPipelineCfg( + soft_contact_margin=SOFT_CONTACT_MARGIN, + ), + ), visualization=visualization_cfg_from_args(args), ) - sim = SimulationManager(config) - - return sim - - -def create_robot(sim: SimulationManager): - """ - Create and configure a robot with an arm and a dexterous hand in the simulation. + return SimulationManager(config) - Args: - sim (SimulationManager): The simulation manager instance. - Returns: - Robot: The configured robot instance added to the simulation. - """ +def create_robot(sim: SimulationManager) -> Robot: + """Add the UR10 and enable particle contact on its pressing flange.""" cfg = URRobotCfg.from_dict( { "robot_type": "ur10", "uid": "UR10", "solver_cfg": {"arm": {"tcp": np.eye(4)}}, + "link_attrs": { + "pressing_flange": { + # ee_link has no collision geometry in the UR10 asset; + # Link6 is the physical flange immediately above it. + "link_names_expr": ["Link6"], + "attrs": { + "collision_props": { + "collision_enabled": True, + "contact_offset": 0.004, + "rest_offset": 0.001, + }, + "material_props": { + "dynamic_friction": SOFT_CONTACT_MU, + }, + "newton_props": { + "collision_props": { + "has_particle_collision": True, + }, + "material_props": { + "ke": SOFT_CONTACT_KE, + "kd": SOFT_CONTACT_KD, + }, + }, + }, + }, + }, "init_qpos": [ 0.0, -np.pi / 2, @@ -111,97 +155,133 @@ def create_robot(sim: SimulationManager): ], } ) - return sim.add_robot(cfg=cfg) + robot = sim.add_robot(cfg=cfg) + if robot is None: + raise RuntimeError("Failed to add the UR10 robot.") + return robot def create_soft_cow(sim: SimulationManager) -> SoftObject: - """create soft cow object in the simulation - - Args: - sim (SimulationManager): The simulation manager instance. - - Returns: - SoftObject: soft cow object - """ - cow: SoftObject = sim.add_soft_object( + """Add the tetrahedral soft cow used by the pressing task.""" + return sim.add_soft_object( cfg=SoftObjectCfg( uid="cow", shape=MeshCfg( fpath=get_resources_data_path("Model", "cow", "cow2.obj"), ), - init_rot=[0, 90, 0], - init_pos=[0.45, -0.1, 0.12], + init_rot=[0.0, 90.0, 0.0], + init_pos=COW_POSITION, + particle_radius=0.005, voxel_attr=SoftbodyVoxelAttributesCfg( - simulation_mesh_resolution=8, - maximal_edge_length=0.5, + triangle_remesh_resolution=24, + simulation_mesh_resolution=16, + voxel_num_relaxation_iters=5, ), physical_attr=SoftbodyPhysicalAttributesCfg( - youngs=5e3, + youngs=5.0e3, poissons=0.45, - density=100, - dynamic_friction=0.1, + density=100.0, + elasticity_damping=0.1, ), - ), + ) ) - return cow - - -def press_cow(sim: SimulationManager, robot: Robot): - """robot press cow softbody with its end link - Args: - sim (SimulationManager): The simulation manager instance. - robot (Robot): The robot instance to be controlled. - """ - start_qpos = robot.get_qpos() - arm_ids = robot.get_joint_ids("arm") - arm_start_qpos = start_qpos[:, arm_ids] - arm_start_xpos = robot.compute_fk(arm_start_qpos, name="arm", to_matrix=True) - press_xpos = arm_start_xpos.clone() - press_xpos[:, :3, 3] = torch.tensor([0.5, -0.1, 0.005], device=press_xpos.device) - - approach_xpos = press_xpos.clone() - approach_xpos[:, 2, 3] += 0.05 - - is_success, approach_qpos = robot.compute_ik( - approach_xpos, joint_seed=arm_start_qpos, name="arm" +def build_press_trajectory(sim: SimulationManager, robot: Robot) -> torch.Tensor: + """Compute the batched approach-and-press trajectory before preparation.""" + arm_joint_names = robot.cfg.control_parts["arm"] + initial_qpos = torch.as_tensor( + robot.cfg.init_qpos, + dtype=torch.float32, + device=sim.device, + ).reshape(1, -1) + initial_qpos = initial_qpos.repeat(sim.num_envs, 1) + + solver_cfg = robot.cfg.solver_cfg["arm"] + solver_cfg.joint_names = list(arm_joint_names) + if solver_cfg.urdf_path is None: + solver_cfg.urdf_path = robot.cfg.fpath + solver = solver_cfg.init_solver(device=sim.device) + + approach_pose = solver.get_fk(initial_qpos) + approach_pose[:, :3, 3] = torch.tensor( + [ + PRESS_POSITION[0], + PRESS_POSITION[1], + PRESS_POSITION[2] + APPROACH_HEIGHT, + ], + dtype=torch.float32, + device=sim.device, + ) + press_pose = approach_pose.clone() + press_pose[:, :3, 3] = torch.tensor( + PRESS_POSITION, + dtype=torch.float32, + device=sim.device, ) - arm_trajectory = torch.concatenate([arm_start_qpos, approach_qpos]) - interp_trajectory = interpolate_with_distance( - trajectory=arm_trajectory[None, :, :], interp_num=50, device=sim.device + approach_success, approach_qpos = solver.get_ik( + target_xpos=approach_pose, + qpos_seed=initial_qpos, + ) + press_success, press_qpos = solver.get_ik( + target_xpos=press_pose, + qpos_seed=approach_qpos, + ) + success = approach_success & press_success + if not bool(torch.all(success)): + failed_envs = torch.nonzero(~success, as_tuple=False).flatten().tolist() + raise RuntimeError(f"IK failed for environment indices {failed_envs}.") + + keyframes = torch.stack( + [initial_qpos, approach_qpos, press_qpos], + dim=1, ) - interp_trajectory = interp_trajectory[0] - for qpos in interp_trajectory: - robot.set_qpos(qpos.unsqueeze(0).repeat(sim.num_envs, 1), joint_ids=arm_ids) - sim.update(step=5) + # Move to the cow over 1.5 seconds, then press by 1.5 cm over 1 second. + return interpolate_with_nums( + trajectory=keyframes, + interp_nums=torch.tensor([150, 100]), + device=sim.device, + ) + +def register_press_trajectory( + sim: SimulationManager, + trajectory: torch.Tensor, +) -> None: + """Register the press after an initial soft-body settling interval.""" + hold = trajectory[:, :1, :].repeat(1, SETTLE_STEPS + 1, 1) + playback = torch.cat([hold, trajectory[:, 1:, :]], dim=1) + sim.register_kinematic_joint_trajectory("UR10", playback) -def main(): - """ - Main function to demonstrate robot simulation. - This function initializes the simulation, creates the robot and other objects, - and performs the press softbody task. - """ +def main() -> None: + """Create the scene, settle the cow, and execute the pressing motion.""" args = parse_arguments() sim = initialize_simulation(args) - robot = create_robot(sim) - soft_cow = create_soft_cow(sim) - sim.prepare() - if not args.headless: - sim.open_window() + try: + robot = create_robot(sim) + create_soft_cow(sim) + trajectory = build_press_trajectory(sim, robot) + register_press_trajectory(sim, trajectory) - press_cow(sim, robot) + sim.prepare() + if not args.headless: + sim.open_window() - logger.log_info("\n Press Ctrl+C to exit simulation loop.") - try: + sim.update(step=SETTLE_STEPS) + if not args.headless: + input("Press Enter to press the soft body...") + sim.update(step=trajectory.shape[1] - 1) + + logger.log_info("\nPress Ctrl+C to exit the simulation loop.") while True: sim.update(step=10) except KeyboardInterrupt: - logger.log_info("\n Exit") + logger.log_info("\nExit") + finally: + sim.destroy() if __name__ == "__main__": diff --git a/examples/sim/demo/softbody_to_cloth.py b/examples/sim/demo/softbody_to_cloth.py new file mode 100644 index 000000000..a194030e9 --- /dev/null +++ b/examples/sim/demo/softbody_to_cloth.py @@ -0,0 +1,422 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Drop a volumetric soft body onto a cloth sheet with Newton VBD.""" + +from __future__ import annotations + +import argparse + +import numpy as np + +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 ( + ClothObjectCfg, + ClothPhysicalAttributesCfg, + LightCfg, + NewtonPhysicsCfg, + RenderCfg, + SoftObjectCfg, + SoftbodyPhysicalAttributesCfg, + SoftbodyVoxelAttributesCfg, +) +from embodichain.lab.sim.material import VisualMaterialCfg +from embodichain.lab.sim.objects import ClothObject, SoftObject +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.visualization import visualization_cfg_from_args +from embodichain.utils import logger + +__all__ = [ + "configure_window_camera", + "create_box_surface_mesh", + "create_cloth", + "create_cloth_grid_mesh", + "create_soft_body", + "initialize_simulation", + "main", + "parse_arguments", + "run_simulation", +] + +FPS = 60 +NUM_SUBSTEPS = 3 +SOLVER_ITERATIONS = 6 +DEFAULT_ITERATIONS = 500 + +CLOTH_SIZE = 2.0 +CLOTH_GRID_CELLS = 28 +CLOTH_POSITION = (-1.0, -1.0, 1.0) +SOFT_BODY_SIZE = (0.6, 0.6, 0.3) +SOFT_BODY_POSITION = (0.0, 0.0, 2.0) + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the coupled deformable demo. + + Returns: + The validated command-line arguments. + """ + parser = argparse.ArgumentParser(description=__doc__) + add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--iterations", + type=int, + default=DEFAULT_ITERATIONS, + help="Number of outer 60 Hz simulation frames.", + ) + parser.set_defaults(device="cuda", physics="newton") + args = parser.parse_args() + if args.physics != "newton": + parser.error("Soft bodies and cloth require --physics newton.") + if not str(args.device).startswith("cuda"): + parser.error("EmbodiChain deformables currently require a CUDA device.") + if args.iterations <= 0: + parser.error("--iterations must be positive.") + return args + + +def initialize_simulation(args: argparse.Namespace) -> SimulationManager: + """Create the Newton VBD simulation manager. + + Args: + args: Parsed command-line arguments. + + Returns: + The configured simulation manager. + """ + cfg = SimulationManagerCfg( + width=1920, + height=1080, + headless=args.headless, + device=args.device, + gpu_id=args.gpu_id, + num_envs=args.num_envs, + arena_space=args.arena_space, + render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=NewtonPhysicsCfg( + physics_dt=1.0 / FPS, + device=args.device, + num_substeps=NUM_SUBSTEPS, + solver_cfg={ + "solver_type": "vbd", + "iterations": SOLVER_ITERATIONS, + "particle_enable_self_contact": True, + "particle_self_contact_radius": 0.01, + "particle_self_contact_margin": 0.02, + "particle_topological_contact_filter_threshold": 3, + "particle_rest_shape_contact_exclusion_radius": 0.05, + "particle_enable_tile_solve": True, + "soft_contact_ke": 1.0e5, + "soft_contact_kd": 1.0e-5, + "soft_contact_mu": 1.0, + }, + ), + visualization=visualization_cfg_from_args(args), + ) + sim = SimulationManager(cfg) + sim.set_emission_light(color=(0.5, 0.5, 0.5), intensity=90.0) + sim.add_light( + LightCfg( + uid="main_light", + intensity=100.0, + radius=10.0, + init_pos=(-4.0, 3.0, 4.5), + ) + ) + return sim + + +def create_box_surface_mesh( + size: tuple[float, float, float], + subdivisions: int = 6, +) -> tuple[np.ndarray, np.ndarray]: + """Create a watertight, subdivided box surface. + + Shared edge and corner vertices keep the mesh suitable for soft-body + tetrahedralization while providing enough render vertices to show its + deformation. + + Args: + size: Box dimensions along the X, Y, and Z axes. + subdivisions: Number of cells along each box edge. + + Returns: + Float32 vertices and int32 outward-facing triangle indices. + + Raises: + ValueError: If the size or subdivision count is invalid. + """ + size_array = np.asarray(size, dtype=np.float32) + if size_array.shape != (3,) or not np.isfinite(size_array).all(): + raise ValueError("size must contain three finite values.") + if np.any(size_array <= 0.0): + raise ValueError("All box dimensions must be positive.") + if subdivisions < 1: + raise ValueError("subdivisions must be at least one.") + + vertices: list[np.ndarray] = [] + triangles: list[tuple[int, int, int]] = [] + vertex_indices: dict[tuple[int, int, int], int] = {} + + def vertex_index(lattice_index: tuple[int, int, int]) -> int: + """Return a shared vertex index for one surface lattice point.""" + index = vertex_indices.get(lattice_index) + if index is not None: + return index + coordinate = ( + np.asarray(lattice_index, dtype=np.float32) / subdivisions - 0.5 + ) * size_array + index = len(vertices) + vertices.append(coordinate) + vertex_indices[lattice_index] = index + return index + + # Each (constant axis, side, U axis, V axis) tuple has U x V pointing + # outward, so both generated triangles have consistent winding. + face_specs = ( + (2, subdivisions, 0, 1), + (2, 0, 1, 0), + (0, subdivisions, 1, 2), + (0, 0, 2, 1), + (1, subdivisions, 2, 0), + (1, 0, 0, 2), + ) + for constant_axis, side, u_axis, v_axis in face_specs: + for v_index in range(subdivisions): + for u_index in range(subdivisions): + corners: list[int] = [] + for u_offset, v_offset in ((0, 0), (1, 0), (1, 1), (0, 1)): + lattice = [0, 0, 0] + lattice[constant_axis] = side + lattice[u_axis] = u_index + u_offset + lattice[v_axis] = v_index + v_offset + corners.append(vertex_index(tuple(lattice))) + triangles.append((corners[0], corners[1], corners[2])) + triangles.append((corners[0], corners[2], corners[3])) + + return ( + np.ascontiguousarray(vertices, dtype=np.float32), + np.ascontiguousarray(triangles, dtype=np.int32), + ) + + +def create_cloth_grid_mesh( + size: float = CLOTH_SIZE, + cells: int = CLOTH_GRID_CELLS, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Create the horizontal cloth grid and its two fixed edge selections. + + Args: + size: Cloth width and depth in metres. + cells: Number of grid cells along each axis. + + Returns: + Vertices, triangles, and fixed-node indices. + + Raises: + ValueError: If the size or cell count is invalid. + """ + if not np.isfinite(size) or size <= 0.0: + raise ValueError("size must be finite and positive.") + if cells < 1: + raise ValueError("cells must be at least one.") + + coordinates = np.linspace(0.0, size, cells + 1, dtype=np.float32) + yy, xx = np.meshgrid(coordinates, coordinates, indexing="ij") + vertices = np.stack( + (xx.reshape(-1), yy.reshape(-1), np.zeros(xx.size, dtype=np.float32)), + axis=1, + ) + + grid_indices = np.arange((cells + 1) ** 2, dtype=np.int32).reshape( + cells + 1, cells + 1 + ) + lower_left = grid_indices[:-1, :-1].reshape(-1) + lower_right = grid_indices[:-1, 1:].reshape(-1) + upper_left = grid_indices[1:, :-1].reshape(-1) + upper_right = grid_indices[1:, 1:].reshape(-1) + triangles = np.concatenate( + ( + np.stack((lower_left, lower_right, upper_right), axis=1), + np.stack((lower_left, upper_right, upper_left), axis=1), + ), + axis=0, + ).astype(np.int32, copy=False) + fixed_indices = np.flatnonzero( + np.isclose(vertices[:, 0], 0.0) | np.isclose(vertices[:, 0], size) + ).astype(np.int32, copy=False) + return ( + np.ascontiguousarray(vertices), + np.ascontiguousarray(triangles), + np.ascontiguousarray(fixed_indices), + ) + + +def create_soft_body(sim: SimulationManager) -> SoftObject: + """Declare the falling soft box with the reference material parameters. + + Args: + sim: Simulation manager that owns the scene declaration. + + Returns: + The declared volume-deformable facade. + """ + vertices, triangles = create_box_surface_mesh(SOFT_BODY_SIZE) + return sim.add_soft_object( + SoftObjectCfg( + uid="falling_soft_body", + shape=MeshCfg( + vertices=vertices, + triangles=triangles, + visual_material=VisualMaterialCfg( + uid="soft_material", + base_color=[0.9, 0.5, 0.2, 1.0], + roughness=0.8, + ), + ), + init_pos=SOFT_BODY_POSITION, + voxel_attr=SoftbodyVoxelAttributesCfg( + simulation_mesh_resolution=10, + ), + physical_attr=SoftbodyPhysicalAttributesCfg( + # This converts exactly to k_mu=8e3 and k_lambda=8e3. + youngs=2.0e4, + poissons=0.25, + density=1.5e2, + elasticity_damping=6.0e-4, + ), + ) + ) + + +def create_cloth(sim: SimulationManager) -> ClothObject: + """Declare the cloth sheet with both X edges fixed. + + Args: + sim: Simulation manager that owns the scene declaration. + + Returns: + The declared surface-deformable facade. + """ + vertices, triangles, fixed_indices = create_cloth_grid_mesh() + particle_flags = np.ones(len(vertices), dtype=np.int32) + particle_flags[fixed_indices] = 0 + return sim.add_cloth_object( + ClothObjectCfg( + uid="cloth_sheet", + shape=MeshCfg( + vertices=vertices, + triangles=triangles, + visual_material=VisualMaterialCfg( + uid="cloth_material", + base_color=[0.4, 0.6, 0.9, 1.0], + roughness=0.8, + ), + ), + init_pos=CLOTH_POSITION, + particle_radius=0.05, + particle_flags=particle_flags, + physical_attr=ClothPhysicalAttributesCfg( + density=5.0e-4, + tri_ke=1.0e5, + tri_ka=1.0e5, + tri_kd=1.0e-5, + edge_ke=0.01, + edge_kd=1.0e-2, + ), + ) + ) + + +def configure_window_camera(sim: SimulationManager) -> None: + """Frame the falling soft body and suspended cloth in the viewer. + + Args: + sim: Prepared simulation manager with an open native window. + """ + window = sim.get_world().get_windows() + if window is not None: + window.set_look_at( + eye=np.asarray([3.0, 3.0, 4.0], dtype=np.float32), + look_at=np.asarray([0.0, 0.0, 1.0], dtype=np.float32), + up=np.asarray([0.0, 0.0, 1.0], dtype=np.float32), + ) + + +def run_simulation( + sim: SimulationManager, + soft_body: SoftObject, + cloth: ClothObject, + iterations: int, +) -> None: + """Advance the coupled scene for a finite number of frames. + + Args: + sim: Prepared simulation manager. + soft_body: Falling volumetric deformable. + cloth: Suspended surface deformable. + iterations: Number of outer simulation frames. + """ + logger.log_info(f"Running soft body to cloth for {iterations} frames at {FPS} Hz.") + logger.log_info( + f"Soft body: {soft_body.get_default_nodal_state().shape[1]} particles, " + f"{soft_body.get_surface_triangles().shape[1]} surface triangles." + ) + logger.log_info( + f"Cloth: {cloth.get_default_nodal_state().shape[1]} particles, " + f"{cloth.get_surface_triangles().shape[1]} triangles." + ) + + for frame in range(iterations): + sim.update(step=1) + if frame % 50 == 0 or frame + 1 == iterations: + soft_height = float( + soft_body.get_current_nodal_position()[..., 2].mean().item() + ) + cloth_height = float( + cloth.get_current_nodal_position()[..., 2].mean().item() + ) + logger.log_info( + f"Frame {frame + 1}/{iterations}, " + f"sim_time={(frame + 1) / FPS:.2f}s, " + f"soft_mean_z={soft_height:.3f}m, " + f"cloth_mean_z={cloth_height:.3f}m" + ) + logger.log_info("Soft body to cloth simulation complete.") + + +def main() -> None: + """Build and run the coupled soft-body/cloth scene.""" + args = parse_arguments() + sim = initialize_simulation(args) + + try: + soft_body = create_soft_body(sim) + cloth = create_cloth(sim) + sim.prepare() + + if not args.headless and sim.open_window(): + configure_window_camera(sim) + run_simulation(sim, soft_body, cloth, args.iterations) + except KeyboardInterrupt: + logger.log_info("\nExit") + finally: + sim.destroy() + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/sim/create_cloth.py b/scripts/tutorials/sim/create_cloth.py index 1e0639fb9..e132f6db6 100644 --- a/scripts/tutorials/sim/create_cloth.py +++ b/scripts/tutorials/sim/create_cloth.py @@ -25,25 +25,31 @@ import os import tempfile import time -import torch + import open3d as o3d -from dexsim.utility.path import get_resources_data_path +import torch + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import ( + MassPropertiesCfg, + NewtonCollisionPipelineCfg, + NewtonPhysicsCfg, RenderCfg, - physics_cfg_for_backend, - RigidObjectCfg, - RigidBodyAttributesCfg, ClothObjectCfg, ClothPhysicalAttributesCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, ) from embodichain.lab.sim.shapes import MeshCfg, CubeCfg from embodichain.lab.sim.objects import ClothObject -def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): +def create_2d_grid_mesh( + width: float, height: float, nx: int = 1, ny: int = 1 +) -> tuple[torch.Tensor, torch.Tensor]: """Create a flat rectangle in the XY plane centered at `origin`. The rectangle is subdivided into an `nx` by `ny` grid (cells) and @@ -59,7 +65,7 @@ def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): # Vectorized vertex positions using PyTorch x_lin = torch.linspace(-w / 2.0, w / 2.0, steps=nx + 1, dtype=torch.float64) y_lin = torch.linspace(-h / 2.0, h / 2.0, steps=ny + 1, dtype=torch.float64) - yy, xx = torch.meshgrid(y_lin, x_lin) # shapes: (ny+1, nx+1) + yy, xx = torch.meshgrid(y_lin, x_lin, indexing="ij") xx_flat = xx.reshape(-1) yy_flat = yy.reshape(-1) zz_flat = torch.full_like(xx_flat, 0, dtype=torch.float64) @@ -77,7 +83,7 @@ def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): return verts, faces -def main(): +def main() -> None: """Main function to create and run the simulation scene.""" # Parse command line arguments @@ -85,18 +91,41 @@ def main(): description="Create a simulation scene with SimulationManager" ) add_env_launcher_args_to_parser(parser) + parser.set_defaults(device="cuda", physics="newton") args = parser.parse_args() + if args.physics != "newton": + parser.error("Cloth requires --physics newton.") + if not str(args.device).startswith("cuda"): + parser.error("Cloth requires a CUDA device.") # Configure the simulation sim_cfg = SimulationManagerCfg( width=1920, height=1080, - headless=True, + headless=args.headless, num_envs=args.num_envs, + arena_space=args.arena_space, + gpu_id=args.gpu_id, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - device="cuda", # soft simulation only supports cuda device + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), - physics_cfg=physics_cfg_for_backend(args.physics), + physics_cfg=NewtonPhysicsCfg( + num_substeps=4, + solver_cfg={ + "solver_type": "vbd", + "iterations": 5, + "particle_enable_self_contact": False, + "particle_self_contact_radius": 0.002, + "particle_self_contact_margin": 0.002, + "particle_enable_tile_solve": True, + "soft_contact_ke": 1.0e4, + "soft_contact_kd": 1.0e-2, + "soft_contact_mu": 0.8, + }, + collision_cfg=NewtonCollisionPipelineCfg( + soft_contact_margin=0.002, + ), + ), visualization=visualization_cfg_from_args(args), ) @@ -117,17 +146,18 @@ def main(): cfg=ClothObjectCfg( uid="cloth", shape=MeshCfg(fpath=cloth_save_path), - init_pos=[0.5, 0.0, 0.3], + init_pos=[0.5, 0.0, 0.8], init_rot=[0, 0, 0], + # The grid spacing is 0.025 m, so avoid Newton's much larger + # 0.1 m default particle radius for this small cloth mesh. + particle_radius=0.01, physical_attr=ClothPhysicalAttributesCfg( - mass=0.01, - youngs=1e9, - poissons=0.4, - thickness=0.004, - bending_stiffness=0.01, - bending_damping=0.1, - dynamic_friction=0.95, - min_position_iters=30, + density=0.02, + tri_ke=2.0e3, + tri_ka=2.0e3, + tri_kd=0.1, + edge_ke=2.0, + edge_kd=0.1, ), ) ) @@ -136,20 +166,20 @@ 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( + mass_props=MassPropertiesCfg(mass=1.0), + material_props=RigidBodyMaterialCfg( + static_friction=0.95, + dynamic_friction=0.9, + restitution=0.01, + ), ), body_type="dynamic", init_pos=[0.5, 0.0, 0.04], init_rot=[0.0, 0.0, 0.0], ) - padding_box = sim.add_rigid_object(cfg=padding_box_cfg) - print("[INFO]: Add soft object complete!") + sim.add_rigid_object(cfg=padding_box_cfg) + print("[INFO]: Add cloth object complete!") sim.prepare() @@ -169,33 +199,29 @@ def run_simulation(sim: SimulationManager, cloth: ClothObject) -> None: Args: sim: The SimulationManager instance to run - soft_obj: soft object + cloth: The cloth object to simulate. """ - step_count = 0 - try: - last_time = time.time() + step_count = 0 + last_time = time.perf_counter() last_step = 0 while True: # Update physics simulation sim.update(step=1) step_count += 1 - # Print FPS every second if step_count % 100 == 0: - current_time = time.time() + current_time = time.perf_counter() elapsed = current_time - last_time fps = ( sim.num_envs * (step_count - last_step) / elapsed - if elapsed > 0 - else 0 + if elapsed > 0.0 + else 0.0 ) print(f"[INFO]: Simulation step: {step_count}, FPS: {fps:.2f}") last_time = current_time last_step = step_count - if step_count % 500 == 0: - cloth.reset() except KeyboardInterrupt: print("\n[INFO]: Stopping simulation...") diff --git a/scripts/tutorials/sim/create_softbody.py b/scripts/tutorials/sim/create_softbody.py index aab5b4112..50db46dda 100644 --- a/scripts/tutorials/sim/create_softbody.py +++ b/scripts/tutorials/sim/create_softbody.py @@ -23,24 +23,24 @@ import argparse import time + from dexsim.utility.path import get_resources_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import ( + NewtonCollisionPipelineCfg, + NewtonPhysicsCfg, RenderCfg, - physics_cfg_for_backend, + SoftObjectCfg, SoftbodyVoxelAttributesCfg, SoftbodyPhysicalAttributesCfg, ) from embodichain.lab.sim.shapes import MeshCfg -from embodichain.lab.sim.objects import ( - SoftObject, - SoftObjectCfg, -) +from embodichain.lab.sim.objects import SoftObject -def main(): +def main() -> None: """Main function to create and run the simulation scene.""" # Parse command line arguments @@ -48,20 +48,44 @@ def main(): description="Create a simulation scene with SimulationManager" ) add_env_launcher_args_to_parser(parser) + parser.set_defaults(device="cuda", physics="newton") args = parser.parse_args() + if args.physics != "newton": + parser.error("Soft bodies require --physics newton.") + if not str(args.device).startswith("cuda"): + parser.error("Soft bodies require a CUDA device.") # Configure the simulation sim_cfg = SimulationManagerCfg( width=1920, height=1080, - headless=True, + headless=args.headless, num_envs=args.num_envs, + arena_space=args.arena_space, + gpu_id=args.gpu_id, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - device="cuda", # soft simulation only supports cuda device + device=args.device, render_cfg=RenderCfg( renderer=args.renderer ), # Enable ray tracing for better visuals - physics_cfg=physics_cfg_for_backend(args.physics), + physics_cfg=NewtonPhysicsCfg( + num_substeps=6, + solver_cfg={ + "solver_type": "vbd", + "iterations": 8, + "particle_enable_self_contact": False, + "particle_self_contact_radius": 0.001, + "particle_self_contact_margin": 0.001, + "particle_topological_contact_filter_threshold": 3, + "particle_enable_tile_solve": True, + "soft_contact_ke": 5.0e4, + "soft_contact_kd": 1.0e-3, + "soft_contact_mu": 1.5, + }, + collision_cfg=NewtonCollisionPipelineCfg( + soft_contact_margin=0.002, + ), + ), visualization=visualization_cfg_from_args(args), ) @@ -77,17 +101,19 @@ def main(): shape=MeshCfg( fpath=get_resources_data_path("Model", "cow", "cow.obj"), ), - init_pos=[0.0, 0.0, 3.0], + init_pos=[0.0, 5.0, 3.0], + particle_radius=0.01, voxel_attr=SoftbodyVoxelAttributesCfg( - simulation_mesh_resolution=8, - maximal_edge_length=0.5, + triangle_remesh_resolution=24, + simulation_mesh_resolution=16, + voxel_num_relaxation_iters=5, ), physical_attr=SoftbodyPhysicalAttributesCfg( - youngs=1e6, - poissons=0.45, - density=100, - dynamic_friction=0.1, - min_position_iters=30, + # Equivalent to the DexSim demo's k_mu=1e4 and k_lambda=5e4. + youngs=2.833333333e4, + poissons=5.0 / 12.0, + density=50.0, + elasticity_damping=2.0e-3, ), ), ) diff --git a/scripts/tutorials/visualization/README.md b/scripts/tutorials/visualization/README.md index 5dbca9bb5..1b95add13 100644 --- a/scripts/tutorials/visualization/README.md +++ b/scripts/tutorials/visualization/README.md @@ -67,8 +67,8 @@ The matching object tutorials can be launched directly: ```bash python scripts/tutorials/sim/create_rigid_object_group.py --viser -python scripts/tutorials/sim/create_softbody.py --viser -python scripts/tutorials/sim/create_cloth.py --viser +python scripts/tutorials/sim/create_softbody.py --physics newton --viser +python scripts/tutorials/sim/create_cloth.py --physics newton --viser ``` The atomic-action tutorials receive the same options through @@ -81,10 +81,10 @@ Application launchers only need to check `--headless` before calling 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 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. +Soft bodies and cloth use the live render surface exposed by their DexSim 0.5 +typed Newton particle-set handles. Volume deformables also retain their +tetrahedral collision-surface topology for physics consumers, while Viser +intentionally publishes the render topology. ## Remote access diff --git a/tests/gym/envs/test_differentiable_embodied_env.py b/tests/gym/envs/test_differentiable_embodied_env.py index 1594b110f..1c932b625 100644 --- a/tests/gym/envs/test_differentiable_embodied_env.py +++ b/tests/gym/envs/test_differentiable_embodied_env.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Tests for DifferentiableEmbodiedEnv.""" +"""Tests for the Newton-only kinematic :class:`DifferentiableEnv`.""" from __future__ import annotations @@ -25,1482 +25,304 @@ import torch import warp as wp -from embodichain.lab.gym.envs.differentiable_env import ( - DifferentiableEmbodiedEnv, -) +import embodichain.lab.gym.envs.differentiable_env as differentiable_env_module +from embodichain.lab.gym.envs.differentiable_env import DifferentiableEnv 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 import NewtonStepFunc, tape_context 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( +def _scale_action_kernel( action: wp.array(dtype=wp.float32), - joint_f: wp.array(dtype=wp.float32), + state: wp.array(dtype=wp.float32), ) -> None: - """Write one tape-tracked action value into Newton joint force.""" - joint_f[0] = action[0] + """Map one action to a task-owned kinematic state.""" + state[0] = 2.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), +def _square_reward_kernel( + state: wp.array(dtype=wp.float32), + reward: 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") - + """Compute a differentiable scalar reward from kinematic state.""" + reward[0] = state[0] * state[0] -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 _bridge_state(*, is_newton_backend: bool = True) -> dict[str, Any]: + """Build a one-dimensional kinematics bridge input on CPU.""" + state_wp = wp.zeros(1, dtype=wp.float32, device="cpu", requires_grad=True) + reward_wp = wp.zeros(1, dtype=wp.float32, device="cpu", requires_grad=True) + manager = SimpleNamespace(is_newton_backend=is_newton_backend) - 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 + def _apply_action(action_wp: Any, tape: Any) -> None: + del tape + wp.launch( + _scale_action_kernel, + dim=1, + inputs=[action_wp, state_wp], + device="cpu", ) - 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 _read_outputs(final_state: Any) -> dict[str, Any]: + assert final_state is state_wp + wp.launch( + _square_reward_kernel, + dim=1, + inputs=[state_wp, reward_wp], + device="cpu", ) - - -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), + "obs": wp.to_torch(state_wp), + "reward": wp.to_torch(reward_wp), "terminated": torch.zeros(1, dtype=torch.bool), "truncated": torch.zeros(1, dtype=torch.bool), "_order": ("obs", "reward", "terminated", "truncated"), - "_grad_track": {}, + "_grad_track": { + "obs": None, + "reward": reward_wp, + "terminated": None, + "truncated": None, + }, } - sim_state: dict[str, Any] = { + return { "manager": manager, - "substeps": _CONTROL_SUBSTEPS, - "physics_dt": nm.physics_dt, - "action_to_control_kernel": action_to_control_kernel, - "kernel_args": ("kernel-argument",), + "action_kernel": _apply_action, + "kernel_args": (), + "step_fn": lambda: state_wp, "obs_reward_fn": _read_outputs, + "last_info": {}, } - 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", + + +def _bare_env() -> DifferentiableEnv: + """Build an uninitialized environment with only its hook dependencies.""" + env = object.__new__(DifferentiableEnv) + env.sim = SimpleNamespace(is_newton_backend=True) + env._apply_action_kernel = lambda _action, tape: None + env._make_kinematic_step_fn = lambda: (lambda: object()) + 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": {}, } - assert [event for event in events if event in tracked_events] == [ - "tape.backward", - "action-gradient.capture", - "tape.reset", - "trajectory.release", - ] + return env def _diff_env_cfg( - requires_grad: bool = True, backend: str = "newton" + requires_grad: bool = True, + backend: str = "newton", ) -> EmbodiedEnvCfg: - if backend == "newton": - physics_cfg = NewtonPhysicsCfg( + """Build the minimum config required for constructor validation.""" + 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", + if backend == "newton" + else DefaultPhysicsCfg() ) - # 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", - ), + return EmbodiedEnvCfg( + sim_cfg=SimulationManagerCfg( + physics_cfg=physics_cfg, + num_envs=2, + headless=True, + ) ) - 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") +def test_public_class_is_renamed_without_legacy_alias() -> None: + """The public module exports only the concise environment name.""" + assert differentiable_env_module.__all__ == ["DifferentiableEnv"] + assert not hasattr(differentiable_env_module, "DifferentiableEmbodiedEnv") - 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" +def test_environment_builds_only_a_kinematic_bridge_state() -> None: + """The base environment exposes no solver mode, substeps, or control hook.""" + env = _bare_env() + expected_state = object() + env._make_kinematic_step_fn = lambda: (lambda: expected_state) - 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", - ), - ) + sim_state = env._build_sim_state_dict(torch.zeros(1)) - 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) + assert sim_state["step_fn"]() is expected_state + assert "step_mode" not in sim_state + assert "substeps" not in sim_state + assert "action_to_control_kernel" not in sim_state + assert "_apply_dynamics_action_kernel" not in DifferentiableEnv.__dict__ + assert "differentiable_step_mode" not in DifferentiableEnv.__dict__ - 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_environment_action_hook_receives_only_action_and_tape() -> None: + """The bridge adapter never supplies a Newton control buffer.""" + env = _bare_env() + action_wp = object() + tape = object() + calls: list[tuple[object, object]] = [] + def _apply_action(action: object, tape: object) -> None: + calls.append((action, tape)) -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) + env._apply_action_kernel = _apply_action - outputs = NewtonStepFunc.apply( - action, - _manager_owned_trajectory_sim_state( - manager, - action_to_control_kernel=lambda _action, *_args: None, - step_mode="dynamics", - ), - ) + sim_state = env._build_sim_state_dict(torch.zeros(1)) + sim_state["action_kernel"](action_wp, tape) - with pytest.raises(RuntimeError, match="injected tape backward failure"): - outputs[0].sum().backward() + assert calls == [(action_wp, tape)] - 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_kinematic_bridge_propagates_reward_gradient_to_action() -> None: + """Warp reverse mode is bridged back to the original torch action.""" + action = torch.tensor([0.5], dtype=torch.float32, requires_grad=True) -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", - ) + _, reward, _, _ = NewtonStepFunc.apply(action, _bridge_state()) + reward.sum().backward() - 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") + assert action.grad is not None + assert torch.allclose(action.grad, torch.tensor([4.0])) - 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 +def test_kinematic_bridge_no_grad_call_releases_tape_synchronously() -> None: + """Inference output does not retain a custom autograd node.""" + action = torch.tensor([0.5], dtype=torch.float32, requires_grad=True) 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 + _, reward, _, _ = NewtonStepFunc.apply(action, _bridge_state()) + assert not reward.requires_grad -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", - ), - ) +def test_kinematic_bridge_rejects_default_backend_before_opening_tape() -> None: + """Direct bridge callers receive the same Newton-only contract.""" + action = torch.tensor([0.5], dtype=torch.float32, requires_grad=True) - 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"): + with pytest.raises(RuntimeError, match="Newton backend"): 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, - ), + action, + _bridge_state(is_newton_backend=False), ) - assert legacy_calls == [] - assert warp.events == [] +def test_kinematic_bridge_requires_a_named_step_callback() -> None: + """An arbitrary dynamics fallback cannot replace the kinematics hook.""" + sim_state = _bridge_state() + sim_state["step_fn"] = None -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(TypeError, match="callable step_fn"): + NewtonStepFunc.apply(torch.zeros(1), sim_state) - 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_tape_context_rejects_default_backend() -> None: + """Expert tape composition remains Newton-only.""" + manager = SimpleNamespace(is_newton_backend=False) + with pytest.raises(RuntimeError, match="Newton backend"): + with tape_context(manager): + pass -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, +def test_kinematic_runtime_does_not_validate_or_expose_a_solver() -> None: + """FK access needs a grad model and live state, not dynamics resources.""" + model = object() + current_state = object() + backend = SimpleNamespace( + model=model, + cfg=SimpleNamespace( + requires_grad=True, + solver_cfg=SimpleNamespace(solver_type="mujoco_warp"), ), + _runtime=SimpleNamespace(current_state=current_state), + state_0=object(), + state_1=object(), ) + runtime = NewtonDifferentiableRuntime(lambda: backend) - 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 runtime.model is model + assert runtime.current_state is current_state + assert runtime.live_states == (backend.state_0, backend.state_1) + assert not hasattr(runtime, "control") + assert not hasattr(runtime, "create_differentiable_trajectory") - assert dynamics_state["step_mode"] == "dynamics" - assert kinematics_state["step_mode"] == "kinematics" +def test_construct_without_requires_grad_raises() -> None: + """Newton kinematic models must opt into Warp gradients.""" + with pytest.raises(RuntimeError, match="requires_grad"): + DifferentiableEnv(_diff_env_cfg(requires_grad=False)) -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": {}, - } +def test_construct_on_default_backend_raises() -> None: + """The Default backend is rejected before environment initialization.""" with pytest.raises( - NotImplementedError, match=r"legacy.*_apply_dynamics_action_kernel" + RuntimeError, + match="DifferentiableEnv requires NewtonPhysicsCfg", ): - 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 == [] + DifferentiableEnv(_diff_env_cfg(backend="default")) -def test_environment_kinematics_hook_keeps_its_strict_legacy_signature( - monkeypatch, +@pytest.mark.parametrize("grad_enabled", (True, False)) +def test_terminal_reset_respects_tape_lifetime( + monkeypatch: pytest.MonkeyPatch, + grad_enabled: bool, ) -> 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 + """Tracked terminal steps defer reset; inference resets synchronously.""" + env = _bare_env() 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": {}, - } + sim_state = {"last_info": {}} + env._build_sim_state_dict = lambda _action: sim_state + + outputs = ( + torch.full((1, 1), 7.0), + torch.full((1,), 3.0), + torch.ones(1, dtype=torch.bool), + torch.zeros(1, dtype=torch.bool), + ) + monkeypatch.setattr( + NewtonStepFunc, + "apply", + staticmethod(lambda _action, _state: outputs), + ) 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")) + context = torch.enable_grad() if grad_enabled else torch.no_grad() + with context: + obs, _, _, _, info = env.step(action) + + if grad_enabled: + assert torch.equal(obs, torch.full((1, 1), 7.0)) + assert reset_calls == [] + assert info["requires_reset_after_backward"] is True + assert torch.equal(info["deferred_reset_ids"], torch.tensor([0])) + else: + assert torch.equal(obs, torch.full((1, 1), -1.0)) + assert len(reset_calls) == 1 + assert "requires_reset_after_backward" not in info 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. - """ + """Import the Franka APG environment after resolving task packages.""" from embodichain_tasks.special.franka_reach_apg import FrankaReachApgEnv return FrankaReachApgEnv def test_franka_kinematics_build_snapshots_live_primal_before_bridge( - monkeypatch, + monkeypatch: pytest.MonkeyPatch, ) -> None: - """Franka must detach taped FK inputs before the parent opens a tape.""" + """Franka detaches taped FK inputs before the parent opens a tape.""" from embodichain_tasks.special import franka_reach_apg env = object.__new__(franka_reach_apg.FrankaReachApgEnv) @@ -1529,11 +351,7 @@ def _parent_build(_self: object, _action: torch.Tensor) -> dict[str, Any]: return {"prepared": True} monkeypatch.setattr(franka_reach_apg.wp, "clone", _clone) - monkeypatch.setattr( - DifferentiableEmbodiedEnv, - "_build_sim_state_dict", - _parent_build, - ) + monkeypatch.setattr(DifferentiableEnv, "_build_sim_state_dict", _parent_build) result = env._build_sim_state_dict(torch.zeros(1, 7)) @@ -1541,24 +359,18 @@ def _parent_build(_self: object, _action: torch.Tensor) -> dict[str, Any]: 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.""" +def test_franka_action_kernel_reads_snapshot_instead_of_live_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The recorded action kernel never captures 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.sim = SimpleNamespace(num_envs=1) env._current_joint_q_snapshot = snapshot_joint_q env._n_joints_per_env = 9 env._wp_device = "cpu" @@ -1581,12 +393,11 @@ def _launch(*_args: Any, inputs: list[object], **_kwargs: Any) -> None: 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, + monkeypatch: pytest.MonkeyPatch, tmp_path, ) -> None: """Detached FK input survives live writes before backward under strict mode.""" @@ -1616,7 +427,7 @@ def test_franka_snapshot_keeps_gradient_after_live_state_mutation_and_matches_fd ) env._action_scale = 0.2 monkeypatch.setattr( - DifferentiableEmbodiedEnv, + DifferentiableEnv, "_build_sim_state_dict", lambda _self, _action: {}, ) @@ -1681,21 +492,20 @@ def _loss(action_value: float) -> float: @pytest.mark.requires_sim @pytest.mark.gpu -def test_franka_apg_smoke_backward(): - """Verify reward is autograd-tracked and action.grad flows back.""" +def test_franka_apg_smoke_backward() -> None: + """Reward remains tracked and produces a finite action gradient.""" try: FrankaReachApgEnv = _import_franka_env() - except FileNotFoundError as e: - pytest.skip(f"Franka URDF not available: {e}") + except FileNotFoundError as exc: + pytest.skip(f"Franka URDF not available: {exc}") 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() + _, reward, _, _, _ = env.step(action) + assert reward.requires_grad + reward.sum().backward() assert action.grad is not None assert torch.isfinite(action.grad).all() finally: @@ -1704,27 +514,26 @@ def test_franka_apg_smoke_backward(): @pytest.mark.requires_sim @pytest.mark.gpu -def test_franka_apg_one_iter_loss_reduces(): - """Verify a single SGD step reduces the APG loss.""" +def test_franka_apg_one_iter_loss_reduces() -> None: + """A short action optimization reduces the kinematic reach loss.""" try: FrankaReachApgEnv = _import_franka_env() - except FileNotFoundError as e: - pytest.skip(f"Franka URDF not available: {e}") + except FileNotFoundError as exc: + pytest.skip(f"Franka URDF not available: {exc}") 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 = [] + optimizer = torch.optim.SGD([action], lr=0.01) + losses: list[float] = [] for _ in range(3): env.reset(seed=0) - opt.zero_grad() + optimizer.zero_grad() _, reward, _, _, _ = env.step(action) loss = (-reward).sum() loss.backward() - opt.step() + optimizer.step() losses.append(loss.detach().item()) assert losses[-1] < losses[0], f"APG did not reduce loss: {losses}" finally: diff --git a/tests/sim/objects/test_cloth_object.py b/tests/sim/objects/test_cloth_object.py index b5238f1a6..f936f6055 100644 --- a/tests/sim/objects/test_cloth_object.py +++ b/tests/sim/objects/test_cloth_object.py @@ -19,7 +19,7 @@ import os from dexsim.utility.path import get_resources_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import ClothPhysicalAttributesCfg +from embodichain.lab.sim.cfg import ClothPhysicalAttributesCfg, NewtonPhysicsCfg from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.objects import ( ClothObject, @@ -31,6 +31,7 @@ import pytest import torch import tempfile +import warp as wp def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): @@ -49,7 +50,7 @@ def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): # Vectorized vertex positions using PyTorch x_lin = torch.linspace(-w / 2.0, w / 2.0, steps=nx + 1, dtype=torch.float64) y_lin = torch.linspace(-h / 2.0, h / 2.0, steps=ny + 1, dtype=torch.float64) - yy, xx = torch.meshgrid(y_lin, x_lin) # shapes: (ny+1, nx+1) + yy, xx = torch.meshgrid(y_lin, x_lin, indexing="ij") xx_flat = xx.reshape(-1) yy_flat = yy.reshape(-1) zz_flat = torch.full_like(xx_flat, 0, dtype=torch.float64) @@ -77,6 +78,9 @@ def setup_simulation(self): device="cuda", num_envs=4, arena_space=3.0, + physics_cfg=NewtonPhysicsCfg( + solver_cfg={"solver_type": "vbd"}, + ), ) # Create the simulation instance @@ -102,14 +106,12 @@ def setup_simulation(self): init_pos=[0.5, 0.0, 0.3], init_rot=[0, 0, 0], physical_attr=ClothPhysicalAttributesCfg( - mass=0.01, - youngs=1e10, - poissons=0.4, - thickness=0.04, - bending_stiffness=0.01, - bending_damping=0.1, - dynamic_friction=0.95, - min_position_iters=30, + density=1.0, + tri_ke=1.0e4, + tri_ka=1.0e4, + tri_kd=10.0, + edge_ke=100.0, + edge_kd=1.0, ), ) ) @@ -164,15 +166,23 @@ def test_unified_deformable_contract(self): 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 + wp.to_torch(particle_set.get_particle_velocities()).clone() + for particle_set in self.cloth.body_data.particle_sets ] ) assert torch.count_nonzero(native_velocities) > 0 torch.testing.assert_close(velocities, native_velocities) + render_vertices = self.cloth.get_surface_vertices() + assert render_vertices.shape[0] == self.sim.num_envs + arena_offsets = torch.as_tensor( + self.sim.arena_offsets, + dtype=render_vertices.dtype, + device=render_vertices.device, + ) + arena_local_vertices = render_vertices - arena_offsets[:, None, :] torch.testing.assert_close( - self.cloth.get_surface_vertices(), - self.cloth.get_current_vertex_position(), + arena_local_vertices, + arena_local_vertices[0].expand_as(arena_local_vertices), ) torch.testing.assert_close( self.cloth.get_surface_triangles(env_ids=[0]), diff --git a/tests/sim/objects/test_deformable_object.py b/tests/sim/objects/test_deformable_object.py index 2791fc076..0a8b08d2b 100644 --- a/tests/sim/objects/test_deformable_object.py +++ b/tests/sim/objects/test_deformable_object.py @@ -19,13 +19,19 @@ from __future__ import annotations from types import SimpleNamespace +from typing import Any, Sequence +import numpy as np +import pytest import torch from embodichain.lab.sim.cfg import ( + ClothPhysicalAttributesCfg, ClothObjectCfg, DeformableObjectCfg, + NewtonPhysicsCfg, SoftObjectCfg, + SoftbodyPhysicalAttributesCfg, SurfaceDeformableObjectCfg, VolumeDeformableObjectCfg, ) @@ -67,6 +73,64 @@ def default_nodal_state_w(self) -> torch.Tensor: return torch.cat((self._pos, torch.zeros_like(self._vel)), dim=-1) +class _ParticleSet: + def __init__(self, offset: float) -> None: + self.positions = torch.tensor( + [[offset, 0.0, 0.0], [offset + 1.0, 0.0, 0.0]], + dtype=torch.float32, + ) + self.velocities = torch.zeros_like(self.positions) + + @property + def particle_count(self) -> int: + return len(self.positions) + + +class _ParticleBatch: + def __init__(self, particle_sets: Sequence[_ParticleSet]) -> None: + self.particle_sets = list(particle_sets) + + def fetch_particle_positions(self, out: torch.Tensor) -> int: + out.copy_(torch.cat([item.positions for item in self.particle_sets])) + return len(self.particle_sets) + + def fetch_particle_velocities(self, out: torch.Tensor) -> int: + out.copy_(torch.cat([item.velocities for item in self.particle_sets])) + return len(self.particle_sets) + + def apply_particle_positions(self, data: torch.Tensor) -> int: + for index, particle_set in enumerate(self.particle_sets): + start = index * particle_set.particle_count + end = start + particle_set.particle_count + particle_set.positions.copy_(data[start:end]) + return len(self.particle_sets) + + def apply_particle_velocities(self, data: torch.Tensor) -> int: + for index, particle_set in enumerate(self.particle_sets): + start = index * particle_set.particle_count + end = start + particle_set.particle_count + particle_set.velocities.copy_(data[start:end]) + return len(self.particle_sets) + + +class _ParticleScene: + def create_particle_set_batch( + self, particle_sets: Sequence[_ParticleSet] + ) -> _ParticleBatch: + return _ParticleBatch(particle_sets) + + +class _RenderParticleSet: + def __init__(self, vertex_count: int) -> None: + self.vertices = np.zeros((vertex_count, 3), dtype=np.float32) + + def get_render_vertices(self) -> np.ndarray: + return self.vertices + + def get_render_triangles(self) -> np.ndarray: + return np.empty((0, 3), dtype=np.int32) + + def test_legacy_configs_specialize_common_deformable_config() -> None: assert issubclass(SoftObjectCfg, VolumeDeformableObjectCfg) assert issubclass(ClothObjectCfg, SurfaceDeformableObjectCfg) @@ -76,6 +140,13 @@ def test_legacy_configs_specialize_common_deformable_config() -> None: assert ClothObjectCfg().deformable_type == "surface" +def test_default_only_deformable_fields_are_not_accepted() -> None: + with pytest.raises(TypeError, match="dynamic_friction"): + SoftbodyPhysicalAttributesCfg(dynamic_friction=0.1) + with pytest.raises(TypeError, match="thickness"): + ClothPhysicalAttributesCfg(thickness=0.01) + + def test_legacy_objects_are_aliases_of_topology_specializations() -> None: assert SoftObject is VolumeDeformableObject assert ClothObject is SurfaceDeformableObject @@ -95,18 +166,112 @@ def test_common_data_contract_combines_and_derives_nodal_state() -> None: 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: +def test_backend_capabilities_are_newton_only() -> 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 + assert not default.supports_volume_deformables + assert not default.supports_surface_deformables + assert not default.supports_soft_bodies + assert not default.supports_cloth + assert newton.supports_volume_deformables + assert newton.supports_surface_deformables + assert newton.supports_soft_bodies + assert newton.supports_cloth + + +def test_manager_rejects_deformables_on_default_backend() -> None: + sim = object.__new__(SimulationManager) + sim.physics = DefaultPhysicsBackend(SimpleNamespace()) + + with pytest.raises(NotImplementedError, match="require the Newton backend"): + sim.add_deformable_object(SoftObjectCfg(uid="soft")) + + +def test_deformable_facade_rejects_default_spawn_scene() -> None: + scene = SimpleNamespace(backend="dexsim") + + with pytest.raises(NotImplementedError, match="Default backend"): + SurfaceDeformableObject( + ClothObjectCfg(uid="cloth"), + entities=[object()], + device=torch.device("cpu"), + spawn_result=scene, + ) + + +@pytest.mark.parametrize("solver_type", ["mujoco_warp", "featherstone"]) +def test_manager_rejects_non_particle_newton_solver(solver_type: str) -> None: + sim = object.__new__(SimulationManager) + sim.physics = NewtonPhysicsBackend(SimpleNamespace()) + sim.physics.solver_type = solver_type + sim.device = torch.device("cuda") + + with pytest.raises(NotImplementedError, match="does not support deformable"): + sim.add_deformable_object(SoftObjectCfg(uid="soft")) + + +def test_manager_rejects_gradient_mode_deformable_mutation() -> None: + sim = object.__new__(SimulationManager) + sim.physics = NewtonPhysicsBackend(SimpleNamespace()) + sim.physics.solver_type = "vbd" + sim.device = torch.device("cuda") + sim.sim_config = SimpleNamespace(physics_cfg=NewtonPhysicsCfg(requires_grad=True)) + + with pytest.raises(NotImplementedError, match="requires_grad=True"): + sim.add_deformable_object(SoftObjectCfg(uid="soft")) + + +def test_particle_data_fetches_and_partially_applies_packed_state() -> None: + particle_sets = [_ParticleSet(0.0), _ParticleSet(10.0)] + data = SurfaceDeformableData( + particle_sets, + _ParticleScene(), + torch.device("cpu"), + ) + + assert data.n_vertices == 2 + torch.testing.assert_close( + data.nodal_pos_w, + torch.stack([item.positions for item in particle_sets]), + ) + torch.testing.assert_close( + data.default_nodal_state_w[..., 3:], + torch.zeros((2, 2, 3)), + ) + + positions = torch.tensor( + [[[20.0, 1.0, 2.0], [21.0, 3.0, 4.0]]], dtype=torch.float32 + ) + velocities = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]], dtype=torch.float32) + data._apply_nodal_state(positions, velocities, env_ids=[1]) + + torch.testing.assert_close( + particle_sets[0].positions[:, 0], torch.tensor([0.0, 1.0]) + ) + torch.testing.assert_close(particle_sets[1].positions, positions[0]) + torch.testing.assert_close(particle_sets[1].velocities, velocities[0]) + + +def test_particle_data_requires_equal_topology() -> None: + particle_sets: list[Any] = [_ParticleSet(0.0), _ParticleSet(1.0)] + particle_sets[1].positions = torch.zeros((3, 3), dtype=torch.float32) + particle_sets[1].velocities = torch.zeros((3, 3), dtype=torch.float32) + + with pytest.raises(ValueError, match="same particle count"): + SurfaceDeformableData( + particle_sets, + _ParticleScene(), + torch.device("cpu"), + ) + + +def test_deformable_rejects_replicated_render_vertex_mismatch() -> None: + deformable = object.__new__(SurfaceDeformableObject) + deformable.device = torch.device("cpu") + + with pytest.raises(RuntimeError, match="render-clone topology mismatch"): + deformable._initialize_topology([_RenderParticleSet(3), _RenderParticleSet(4)]) def test_manager_generic_and_legacy_getters_share_one_registry() -> None: diff --git a/tests/sim/objects/test_soft_object.py b/tests/sim/objects/test_soft_object.py index 7fe8352a9..4ffa28b2f 100644 --- a/tests/sim/objects/test_soft_object.py +++ b/tests/sim/objects/test_soft_object.py @@ -20,14 +20,13 @@ from dexsim.utility.path import get_resources_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( - RenderCfg, + NewtonPhysicsCfg, SoftbodyVoxelAttributesCfg, SoftbodyPhysicalAttributesCfg, ) from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.objects import ( DeformableObject, - SoftBodyData, SoftObject, SoftObjectCfg, VolumeDeformableObject, @@ -38,17 +37,6 @@ COW_PATH = get_resources_data_path("Model", "cow", "cow.obj") -def test_degenerate_soft_body_surface_is_empty() -> None: - """Degenerate collision geometry does not prevent visualization startup.""" - data = object.__new__(SoftBodyData) - data.device = torch.device("cpu") - data._rest_position_buffer = torch.zeros((1, 3, 4), dtype=torch.float32) - - triangles = data.collision_surface_triangles - - assert triangles.shape == (0, 3) - - class BaseSoftObjectTest: def setup_simulation(self): sim_cfg = SimulationManagerCfg( @@ -57,8 +45,13 @@ def setup_simulation(self): headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) device="cuda", - num_envs=4, + # DexSim 0.5 currently changes the cow render topology while + # cloning it; keep the functional volume test single-instance. + num_envs=1, arena_space=3.0, + physics_cfg=NewtonPhysicsCfg( + solver_cfg={"solver_type": "vbd"}, + ), ) # Create the simulation instance @@ -67,7 +60,7 @@ def setup_simulation(self): assert os.path.isfile(COW_PATH) # Enable manual physics update for precise control - self.num_envs = 4 + self.num_envs = 1 # add softbody to the scene self.cow: SoftObject = self.sim.add_soft_object( @@ -79,14 +72,12 @@ def setup_simulation(self): init_pos=[0.0, 0.0, 3.0], voxel_attr=SoftbodyVoxelAttributesCfg( simulation_mesh_resolution=8, - maximal_edge_length=0.5, ), physical_attr=SoftbodyPhysicalAttributesCfg( youngs=1e6, poissons=0.45, density=100, - dynamic_friction=0.1, - min_position_iters=30, + elasticity_damping=0.1, ), ), ) @@ -124,14 +115,10 @@ def test_unified_deformable_contract(self): 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]), - ) + render_vertices = self.cow.get_surface_vertices() + render_triangles = self.cow.get_surface_triangles(env_ids=[0]) + assert render_vertices.shape[0] == self.sim.num_envs + assert int(render_triangles.max()) < render_vertices.shape[1] def test_remove(self): with pytest.raises(NotImplementedError, match="pending removal"): diff --git a/tests/sim/spawn/test_descriptors.py b/tests/sim/spawn/test_descriptors.py index 822b22734..d75f747de 100644 --- a/tests/sim/spawn/test_descriptors.py +++ b/tests/sim/spawn/test_descriptors.py @@ -34,10 +34,8 @@ CollisionDesc, CollisionApproximation, DexsimCollisionDesc, - DexsimClothPhysicsDesc, DexsimJointDesc, DexsimPhysicsDesc, - DexsimSoftBodyPhysicsDesc, JointDesc, LinkDesc, NewtonCollisionDesc, @@ -96,29 +94,31 @@ DEFORMABLE_MESH_PATH = "/assets/deformable.obj" -def test_soft_descriptor_projects_current_dexsim_particle_schema() -> None: +def test_soft_descriptor_uses_newton_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), + particle_radius=0.02, + particle_flags=0, + validate_mesh=True, voxel_attr=SoftbodyVoxelAttributesCfg( - triangle_remesh_resolution=remesh_resolution, - triangle_simplify_target=simplify_target, - simulation_mesh_resolution=voxel_resolution, + triangle_remesh_resolution=12, + triangle_simplify_target=40, + simulation_mesh_resolution=16, + voxel_num_relaxation_iters=7, + voxel_rel_min_tet_volume=0.08, + voxel_surface_dist_ratio=0.3, + embedding_impl="dexsim_exact_cpu", ), physical_attr=SoftbodyPhysicalAttributesCfg( youngs=youngs, poissons=poissons, - density=density, - dynamic_friction=dynamic_friction, - min_position_iters=min_position_iters, + density=75.0, + elasticity_damping=0.2, + surface_tri_ke=1.0, + surface_edge_ke=2.0, ), ) @@ -126,35 +126,46 @@ def test_soft_descriptor_projects_current_dexsim_particle_schema() -> None: assert isinstance(descriptor, SoftBodyDesc) assert descriptor.mesh.file_path == DEFORMABLE_MESH_PATH + assert descriptor.particle_radius == pytest.approx(0.02) + assert descriptor.particle_flags == 0 + assert descriptor.validate_mesh is True 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.volume_density == pytest.approx(75.0) 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 descriptor.physics.k_damp == pytest.approx(0.2) + assert descriptor.physics.surface_tri_ke == pytest.approx(1.0) + assert descriptor.physics.surface_edge_ke == pytest.approx(2.0) + assert descriptor.physics.dexsim is None + assert descriptor.meshing.proxy_simplify_target == 40 + assert descriptor.meshing.proxy_remesh_resolution == 12 + assert descriptor.meshing.voxel_resolution == 16 + assert descriptor.meshing.voxel_num_relaxation_iters == 7 + assert descriptor.meshing.voxel_rel_min_tet_volume == pytest.approx(0.08) + assert descriptor.meshing.voxel_surface_dist_ratio == pytest.approx(0.3) + assert descriptor.meshing.embedding_impl == "dexsim_exact_cpu" 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 +def test_cloth_descriptor_uses_newton_particle_schema() -> None: cfg = ClothObjectCfg( uid="cloth", shape=MeshCfg(fpath=DEFORMABLE_MESH_PATH), + particle_radius=0.01, + particle_flags=np.asarray([0, 1, 0], dtype=np.int32), + validate_mesh=True, physical_attr=ClothPhysicalAttributesCfg( - density=density, - mass=mass, - thickness=thickness, - bending_stiffness=bending_stiffness, + density=2.5, + tri_ke=100.0, + tri_ka=90.0, + tri_kd=5.0, + edge_ke=20.0, + edge_kd=2.0, + add_springs=True, + spring_ke=30.0, + spring_kd=3.0, ), ) @@ -162,15 +173,159 @@ def test_cloth_descriptor_projects_current_dexsim_particle_schema() -> None: assert isinstance(descriptor, ClothDesc) assert descriptor.mesh.file_path == DEFORMABLE_MESH_PATH + assert descriptor.particle_radius == pytest.approx(0.01) + np.testing.assert_array_equal(descriptor.particle_flags, [0, 1, 0]) + assert descriptor.validate_mesh is True 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 descriptor.physics.surface_density == pytest.approx(2.5) + assert descriptor.physics.tri_ke == pytest.approx(100.0) + assert descriptor.physics.tri_ka == pytest.approx(90.0) + assert descriptor.physics.tri_kd == pytest.approx(5.0) + assert descriptor.physics.edge_ke == pytest.approx(20.0) + assert descriptor.physics.edge_kd == pytest.approx(2.0) + assert descriptor.physics.add_springs is True + assert descriptor.physics.spring_ke == pytest.approx(30.0) + assert descriptor.physics.spring_kd == pytest.approx(3.0) + assert descriptor.physics.dexsim is None + assert materials == {} + + +def test_cloth_descriptor_preserves_array_mesh_vertex_order() -> None: + vertices = np.asarray( + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + dtype=np.float32, + ) + triangles = np.asarray([[2, 0, 1]], dtype=np.int32) + uv_coords = np.asarray( + [[0.0, 0.0], [1.0, 0.0], [0.5, 1.0]], + dtype=np.float32, + ) + cfg = ClothObjectCfg( + uid="cloth", + shape=MeshCfg( + vertices=vertices, + triangles=triangles, + uv_coords=uv_coords, + ), + particle_flags=[0, 1, 1], + ) + + cfg.validate() + descriptor, _ = cloth_desc_from_cfg(cfg) + + assert descriptor.mesh.file_path is None + np.testing.assert_array_equal(descriptor.mesh.vertices, vertices) + np.testing.assert_array_equal(descriptor.mesh.triangles, triangles) + np.testing.assert_array_equal(descriptor.mesh.uv_coords, uv_coords) + np.testing.assert_array_equal(descriptor.particle_flags, [0, 1, 1]) + + +def test_cloth_descriptor_supports_independent_visual_mesh() -> None: + visual_mesh_path = "/assets/deformable_visual.obj" + cfg = ClothObjectCfg( + uid="cloth", + shape=MeshCfg( + vertices=np.asarray( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + dtype=np.float32, + ), + triangles=np.asarray([[0, 1, 2]], dtype=np.int32), + ), + visual_shape=MeshCfg(fpath=visual_mesh_path), + visual_binding_mode="nearest_vertex", + ) + + descriptor, materials = cloth_desc_from_cfg(cfg) + + assert descriptor.mesh.file_path is None + assert descriptor.visual_mesh is not None + assert descriptor.visual_mesh.file_path == visual_mesh_path + assert descriptor.visual_binding_mode == "nearest_vertex" assert materials == {} +def test_cloth_descriptor_rejects_unknown_visual_binding_mode() -> None: + cfg = ClothObjectCfg( + uid="cloth", + shape=MeshCfg(fpath=DEFORMABLE_MESH_PATH), + visual_binding_mode="unsupported", + ) + + with pytest.raises(ValueError, match="visual_binding_mode"): + cloth_desc_from_cfg(cfg) + + +def test_cloth_descriptor_rejects_multiple_mesh_sources() -> None: + cfg = ClothObjectCfg( + uid="cloth", + shape=MeshCfg( + fpath=DEFORMABLE_MESH_PATH, + vertices=np.zeros((3, 3), dtype=np.float32), + triangles=np.asarray([[0, 1, 2]], dtype=np.int32), + ), + ) + + with pytest.raises(ValueError, match="either fpath or vertices/triangles"): + cloth_desc_from_cfg(cfg) + + +def test_cloth_descriptor_rejects_missing_mesh_source_after_config_validation() -> None: + cfg = ClothObjectCfg(uid="cloth", shape=MeshCfg()) + + cfg.validate() + with pytest.raises(ValueError, match="non-empty fpath or vertices/triangles"): + cloth_desc_from_cfg(cfg) + + +def test_soft_descriptor_rejects_invalid_poisson_ratio() -> None: + cfg = SoftObjectCfg( + uid="soft", + shape=MeshCfg(fpath=DEFORMABLE_MESH_PATH), + physical_attr=SoftbodyPhysicalAttributesCfg(poissons=0.5), + ) + + with pytest.raises(ValueError, match="poissons"): + soft_desc_from_cfg(cfg) + + +@pytest.mark.parametrize("particle_radius", [0.0, float("nan")]) +def test_cloth_descriptor_rejects_invalid_particle_radius( + particle_radius: float, +) -> None: + cfg = ClothObjectCfg( + uid="cloth", + shape=MeshCfg(fpath=DEFORMABLE_MESH_PATH), + particle_radius=particle_radius, + ) + + with pytest.raises(ValueError, match="particle_radius"): + cloth_desc_from_cfg(cfg) + + +@pytest.mark.parametrize( + "particle_flags", + [ + True, + np.iinfo(np.int32).max + 1, + [0.0, 1.0], + [-1, 1], + [np.iinfo(np.int32).max + 1], + np.zeros((1, 2), dtype=np.int32), + ], +) +def test_cloth_descriptor_rejects_invalid_particle_flags( + particle_flags: object, +) -> None: + cfg = ClothObjectCfg( + uid="cloth", + shape=MeshCfg(fpath=DEFORMABLE_MESH_PATH), + particle_flags=particle_flags, + ) + + with pytest.raises((TypeError, ValueError), match="particle_flags"): + cloth_desc_from_cfg(cfg) + + def _resolved_articulation_desc() -> ArticulationDesc: source_inertia = np.ones(3, dtype=np.float32) base = LinkDesc( @@ -707,6 +862,25 @@ def test_mesh_cfg_collision_fields_remain_compatibility_fallbacks() -> None: assert descriptor.collisions[0].decomp_max_hulls == 3 +def test_rigid_descriptor_preserves_array_mesh_data() -> None: + vertices = np.asarray( + [[2.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 2.0]], + dtype=np.float32, + ) + triangles = np.asarray([[2, 0, 1]], dtype=np.int32) + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg(vertices=vertices, triangles=triangles), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + np.testing.assert_array_equal(descriptor.renders[0].vertices, vertices) + np.testing.assert_array_equal(descriptor.renders[0].triangles, triangles) + np.testing.assert_array_equal(descriptor.collisions[0].vertices, vertices) + np.testing.assert_array_equal(descriptor.collisions[0].triangles, triangles) + + def test_backend_blocks_reject_portable_fields() -> None: cfg = RigidObjectCfg( uid="cube", @@ -1524,9 +1698,14 @@ def test_spawn_post_config_only_applies_render_uv() -> None: render_body = Mock() entity = Mock() entity.get_render_body.return_value = render_body + entity.joint_dof_layout = [] articulation = object.__new__(Articulation) articulation.cfg = SimpleNamespace(compute_uv=True) articulation._entities = [entity] + articulation._mimic_info = SimpleNamespace( + mimic_id=np.asarray([], dtype=np.int32), + mimic_parent=np.asarray([], dtype=np.int32), + ) articulation.__dict__["link_names"] = ["base"] articulation._set_default_joint_drive = Mock() @@ -1539,7 +1718,10 @@ def test_spawn_post_config_only_applies_render_uv() -> None: def test_spawn_post_config_applies_default_only_root_properties() -> None: native_articulation = Mock() - entity = SimpleNamespace(_physics_binding=native_articulation) + entity = SimpleNamespace( + _physics_binding=native_articulation, + joint_dof_layout=[], + ) articulation = object.__new__(Articulation) articulation.cfg = ArticulationCfg( articulation_props=ArticulationRootPropertiesCfg( @@ -1548,8 +1730,16 @@ def test_spawn_post_config_applies_default_only_root_properties() -> None: min_velocity_iters=2, ) ) - articulation._spawn_result = SimpleNamespace(backend="dexsim") + articulation._spawn_result = SimpleNamespace( + backend="dexsim", + topology_revision=1, + ) articulation._entities = [entity] + articulation._prepared_default_root_topology_revision = -1 + articulation._mimic_info = SimpleNamespace( + mimic_id=np.asarray([], dtype=np.int32), + mimic_parent=np.asarray([], dtype=np.int32), + ) articulation._apply_spawn_config() @@ -1562,7 +1752,10 @@ def test_spawn_post_config_applies_default_only_root_properties() -> None: def test_newton_skips_default_only_articulation_root_properties() -> None: native_articulation = Mock() - entity = SimpleNamespace(_physics_binding=native_articulation) + entity = SimpleNamespace( + _physics_binding=native_articulation, + joint_dof_layout=[], + ) articulation = object.__new__(Articulation) articulation.cfg = ArticulationCfg( articulation_props=ArticulationRootPropertiesCfg( @@ -1573,6 +1766,10 @@ def test_newton_skips_default_only_articulation_root_properties() -> None: ) articulation._spawn_result = SimpleNamespace(backend="newton") articulation._entities = [entity] + articulation._mimic_info = SimpleNamespace( + mimic_id=np.asarray([], dtype=np.int32), + mimic_parent=np.asarray([], dtype=np.int32), + ) articulation._apply_spawn_config() diff --git a/tests/sim/test_backend_parity.py b/tests/sim/test_backend_parity.py index c09f2149a..c98ca5d25 100644 --- a/tests/sim/test_backend_parity.py +++ b/tests/sim/test_backend_parity.py @@ -46,10 +46,10 @@ # 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}, + "volume_deformables": {"default": False, "newton": True}, + "surface_deformables": {"default": False, "newton": True}, + "soft_bodies": {"default": False, "newton": True}, + "cloth": {"default": False, "newton": True}, "rigid_object_group": {"default": True, "newton": True}, "can_disable_manual_update": {"default": True, "newton": False}, } diff --git a/tests/sim/test_differentiable_stepper.py b/tests/sim/test_differentiable_stepper.py deleted file mode 100644 index 4006da542..000000000 --- a/tests/sim/test_differentiable_stepper.py +++ /dev/null @@ -1,100 +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. -# ---------------------------------------------------------------------------- -"""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_runtime_controls.py b/tests/sim/test_runtime_controls.py new file mode 100644 index 000000000..531fba080 --- /dev/null +++ b/tests/sim/test_runtime_controls.py @@ -0,0 +1,137 @@ +# ---------------------------------------------------------------------------- +# 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 manager-owned Newton runtime-control adapters.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest + +from embodichain.lab.sim._runtime_controls import ( + _KinematicNodalTrajectoryControl, +) + +pytestmark = pytest.mark.no_sim + + +class _ArrayView: + """Host-array stand-in for a borrowed Warp particle view.""" + + def __init__(self, values: np.ndarray) -> None: + self.values = values + + def numpy(self) -> np.ndarray: + """Return a host snapshot matching Warp's ``numpy`` method.""" + return self.values.copy() + + +class _ParticleSet: + """Minimal particle-set facade used by the runtime-control tests.""" + + def __init__(self) -> None: + self.positions = np.asarray( + [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]], + dtype=np.float32, + ) + self.fixed_indices: np.ndarray | None = None + + @property + def particle_count(self) -> int: + """Return the number of test particles.""" + return len(self.positions) + + def get_particle_positions(self) -> _ArrayView: + """Return the current particle positions.""" + return _ArrayView(self.positions) + + def set_particle_positions(self, positions: np.ndarray) -> None: + """Store one complete position snapshot.""" + self.positions = np.asarray(positions, dtype=np.float32).copy() + + def fix_particles(self, particle_indices: np.ndarray) -> None: + """Record which particles were made kinematic.""" + self.fixed_indices = np.asarray(particle_indices, dtype=np.int32).copy() + + +def test_kinematic_nodal_control_interpolates_offsets_per_substep() -> None: + particle_set = _ParticleSet() + solver = SimpleNamespace(rebuild_bvh=MagicMock()) + current_state = object() + context = SimpleNamespace( + result=SimpleNamespace(get_particle_set=lambda _target: particle_set), + solver=solver, + current_state=current_state, + ) + offsets = np.asarray( + [ + [[0.0, 0.0, 0.0]], + [[0.0, 2.0, 0.0]], + ], + dtype=np.float32, + ) + control = _KinematicNodalTrajectoryControl( + "arena_0/cloth", + np.asarray([1], dtype=np.int32), + offsets, + fps=10.0, + rebuild_self_contact_bvh=True, + ) + + control.initialize(context) + control(context, substep_index=0, substep_count=2, substep_dt=0.05) + np.testing.assert_allclose(particle_set.positions[1], [2.0, 0.0, 0.0]) + control(context, substep_index=1, substep_count=2, substep_dt=0.05) + + assert particle_set.fixed_indices is None + np.testing.assert_allclose(particle_set.positions[0], [0.0, 0.0, 0.0]) + np.testing.assert_allclose(particle_set.positions[1], [2.0, 1.0, 0.0]) + solver.rebuild_bvh.assert_called_once_with(current_state) + assert control.exclusive_resource_claims() == ( + ("kinematic_nodal_trajectory", "arena_0/cloth"), + ) + + +def test_kinematic_nodal_control_holds_last_unrated_sample() -> None: + particle_set = _ParticleSet() + context = SimpleNamespace( + result=SimpleNamespace(get_particle_set=lambda _target: particle_set), + solver=object(), + current_state=object(), + ) + offsets = np.asarray( + [ + [[0.0, 0.0, 0.0]], + [[1.0, 0.0, 0.0]], + ], + dtype=np.float32, + ) + control = _KinematicNodalTrajectoryControl( + "arena_0/cloth", + np.asarray([1], dtype=np.int32), + offsets, + fps=None, + rebuild_self_contact_bvh=False, + ) + control.initialize(context) + + for _ in range(3): + control(context, substep_index=0, substep_count=1, substep_dt=0.01) + + np.testing.assert_allclose(particle_set.positions[1], [3.0, 0.0, 0.0]) diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 07a7bbcf7..68c26a06d 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -28,6 +28,7 @@ import embodichain.lab.sim.sim_manager as sim_manager_module from embodichain.lab.sim.cfg import ( + ClothObjectCfg, DefaultPhysicsCfg, RobotCfg, RobotPresetCfg, @@ -237,6 +238,30 @@ def _make_visualization_sim_manager() -> ( return sim, runtime +def _make_runtime_control_sim_manager( + *, + num_envs: int = 2, + backend: str = "newton", +) -> tuple[SimulationManager, MagicMock]: + """Create a manager stub at the pre-prepare runtime-control boundary.""" + sim = object.__new__(SimulationManager) + sim.sim_config = SimpleNamespace(num_envs=num_envs) + sim.physics = SimpleNamespace(name=backend) + sim._robots = {"robot": object()} + sim._articulations = {} + sim._rigid_objects = {"table": object()} + sim._deformable_objects = { + "cloth": SimpleNamespace( + cfg=ClothObjectCfg(uid="cloth", particle_flags=[0, 1, 0]) + ) + } + spawn_scene = MagicMock() + spawn_scene.arena_names = tuple(f"arena_{index}" for index in range(num_envs)) + spawn_scene.builder.is_finalized = False + sim._spawn_scene = spawn_scene + return sim, spawn_scene.builder + + def test_flush_cleanup_queue_returns_immediately_when_no_destroy_is_pending( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -606,6 +631,207 @@ def start_visualization(sim: SimulationManager) -> None: assert sim._arenas == [] +def test_register_kinematic_joint_trajectory_expands_each_arena() -> None: + sim, builder = _make_runtime_control_sim_manager() + frame_count = 3 + dof_count = 2 + positions = torch.arange( + sim.num_envs * frame_count * dof_count, + dtype=torch.float32, + ).reshape(sim.num_envs, frame_count, dof_count) + root_poses = np.tile( + np.eye(4, dtype=np.float32), + (sim.num_envs, frame_count, 1, 1), + ) + + sim.register_kinematic_joint_trajectory( + "robot", + positions, + fps=50.0, + root_poses=root_poses, + ) + + assert builder.add_runtime_control.call_count == sim.num_envs + for env_index, control_call in enumerate( + builder.add_runtime_control.call_args_list + ): + control = control_call.args[0] + assert control.target == f"arena_{env_index}/robot" + assert control.fps == 50.0 + np.testing.assert_array_equal( + control.joint_positions, + positions[env_index].numpy(), + ) + np.testing.assert_array_equal( + control.root_poses, + root_poses[env_index], + ) + + +def test_register_kinematic_joint_trajectory_rejects_non_newton_backend() -> None: + sim, builder = _make_runtime_control_sim_manager(backend="default") + positions = np.zeros((sim.num_envs, 2, 1), dtype=np.float32) + + with pytest.raises(RuntimeError, match="require the Newton backend"): + sim.register_kinematic_joint_trajectory("robot", positions) + + builder.add_runtime_control.assert_not_called() + + +def test_register_kinematic_joint_trajectory_rejects_unknown_asset() -> None: + sim, builder = _make_runtime_control_sim_manager() + positions = np.zeros((sim.num_envs, 2, 1), dtype=np.float32) + + with pytest.raises(KeyError, match="missing"): + sim.register_kinematic_joint_trajectory("missing", positions) + + builder.add_runtime_control.assert_not_called() + + +def test_register_kinematic_joint_trajectory_rejects_wrong_arena_batch() -> None: + sim, builder = _make_runtime_control_sim_manager(num_envs=2) + positions = np.zeros((1, 2, 1), dtype=np.float32) + + with pytest.raises(ValueError, match=r"\(2, frames, dof\)"): + sim.register_kinematic_joint_trajectory("robot", positions) + + builder.add_runtime_control.assert_not_called() + + +def test_register_kinematic_joint_trajectory_rejects_finalized_scene() -> None: + sim, builder = _make_runtime_control_sim_manager() + builder.is_finalized = True + positions = np.zeros((sim.num_envs, 2, 1), dtype=np.float32) + + with pytest.raises(RuntimeError, match=r"before SimulationManager\.prepare"): + sim.register_kinematic_joint_trajectory("robot", positions) + + builder.add_runtime_control.assert_not_called() + + +def test_register_contact_material_schedule_expands_each_arena() -> None: + sim, builder = _make_runtime_control_sim_manager() + friction_track = ((0.0, 0.5), (1.0, 0.1)) + + sim.register_contact_material_schedule( + "robot", + {"dynamic_friction": friction_track}, + link_names=("left_finger", "right_finger"), + ) + + assert builder.add_runtime_control.call_count == sim.num_envs + for env_index, control_call in enumerate( + builder.add_runtime_control.call_args_list + ): + control = control_call.args[0] + assert control.target == f"arena_{env_index}/robot" + assert control.link_names == ("left_finger", "right_finger") + assert control.property_names == ("dynamic_friction",) + np.testing.assert_allclose(control.keyframe_times, (0.0, 1.0)) + np.testing.assert_allclose(control.keyframe_values, (0.5, 0.1)) + + +def test_register_contact_material_schedule_rejects_unknown_asset() -> None: + sim, builder = _make_runtime_control_sim_manager() + + with pytest.raises(KeyError, match="missing"): + sim.register_contact_material_schedule( + "missing", + {"dynamic_friction": ((0.0, 0.5),)}, + ) + + builder.add_runtime_control.assert_not_called() + + +def test_register_particle_contact_material_schedule_adds_global_control() -> None: + sim, builder = _make_runtime_control_sim_manager() + friction_track = ((0.0, 0.5), (1.0, 1.2)) + + sim.register_particle_contact_material_schedule( + {"dynamic_friction": friction_track} + ) + + builder.add_runtime_control.assert_called_once() + control = builder.add_runtime_control.call_args.args[0] + np.testing.assert_allclose( + control.tracks["dynamic_friction"], + friction_track, + ) + + +def test_contact_material_schedules_reject_finalized_scene() -> None: + sim, builder = _make_runtime_control_sim_manager() + builder.is_finalized = True + keyframes = {"dynamic_friction": ((0.0, 0.5),)} + + with pytest.raises(RuntimeError, match=r"before SimulationManager\.prepare"): + sim.register_contact_material_schedule("table", keyframes) + with pytest.raises(RuntimeError, match=r"before SimulationManager\.prepare"): + sim.register_particle_contact_material_schedule(keyframes) + + builder.add_runtime_control.assert_not_called() + + +def test_register_kinematic_nodal_trajectory_expands_each_arena() -> None: + sim, builder = _make_runtime_control_sim_manager() + node_indices = np.asarray([0, 2], dtype=np.int32) + sample_count = 3 + offsets = np.arange( + sim.num_envs * sample_count * len(node_indices) * 3, + dtype=np.float32, + ).reshape(sim.num_envs, sample_count, len(node_indices), 3) + + sim.register_kinematic_nodal_trajectory( + "cloth", + node_indices, + offsets, + fps=60.0, + rebuild_self_contact_bvh=True, + ) + + assert builder.add_runtime_control.call_count == sim.num_envs + for env_index, control_call in enumerate( + builder.add_runtime_control.call_args_list + ): + control = control_call.args[0] + assert control.target == f"arena_{env_index}/cloth" + assert control.fps == pytest.approx(60.0) + assert control.rebuild_self_contact_bvh is True + np.testing.assert_array_equal(control.node_indices, node_indices) + np.testing.assert_array_equal(control.position_offsets, offsets[env_index]) + + +def test_register_kinematic_nodal_trajectory_rejects_active_nodes() -> None: + sim, builder = _make_runtime_control_sim_manager() + offsets = np.zeros((sim.num_envs, 2, 1, 3), dtype=np.float32) + + with pytest.raises(ValueError, match="ACTIVE particle flag"): + sim.register_kinematic_nodal_trajectory("cloth", [1], offsets) + + builder.add_runtime_control.assert_not_called() + + +def test_register_kinematic_nodal_trajectory_rejects_wrong_shape() -> None: + sim, builder = _make_runtime_control_sim_manager() + offsets = np.zeros((sim.num_envs, 2, 3), dtype=np.float32) + + with pytest.raises(ValueError, match="position_offsets"): + sim.register_kinematic_nodal_trajectory("cloth", [0], offsets) + + builder.add_runtime_control.assert_not_called() + + +def test_register_kinematic_nodal_trajectory_rejects_finalized_scene() -> None: + sim, builder = _make_runtime_control_sim_manager() + builder.is_finalized = True + offsets = np.zeros((sim.num_envs, 2, 1, 3), dtype=np.float32) + + with pytest.raises(RuntimeError, match=r"before SimulationManager\.prepare"): + sim.register_kinematic_nodal_trajectory("cloth", [0], offsets) + + builder.add_runtime_control.assert_not_called() + + def test_add_robot_resolves_backend_preset_before_declaration() -> None: @configclass class TestRobotPresetCfg(RobotPresetCfg): diff --git a/tests/sim/test_sim_manager_cfg.py b/tests/sim/test_sim_manager_cfg.py index f7257a97d..43bd6a174 100644 --- a/tests/sim/test_sim_manager_cfg.py +++ b/tests/sim/test_sim_manager_cfg.py @@ -297,6 +297,17 @@ def test_newton_physics_cfg_converts_mapping_solver_cfg_to_dexsim_cfg() -> None: assert dexsim_cfg.solver_cfg.use_mujoco_contacts is False +def test_newton_physics_cfg_accepts_mjvbd_solver_alias() -> None: + from dexsim.engine.newton_physics import MJVBDSolverCfg + + cfg = NewtonPhysicsCfg(solver_cfg={"class_type": "MJVBDSolverCfg"}) + + dexsim_cfg = cfg.to_dexsim_cfg(gpu_id=0) + + assert isinstance(dexsim_cfg.solver_cfg, MJVBDSolverCfg) + assert dexsim_cfg.solver_cfg.solver_type == "mjvbd" + + def test_newton_physics_cfg_directly_accepts_dexsim_solver_cfg_object() -> None: from dexsim.engine.newton_physics import XPBDSolverCfg