diff --git a/.github/workflows/aws_gpu_benchmarks.yml b/.github/workflows/aws_gpu_benchmarks.yml index 219b5b4e3e..8a9fca8add 100644 --- a/.github/workflows/aws_gpu_benchmarks.yml +++ b/.github/workflows/aws_gpu_benchmarks.yml @@ -112,6 +112,7 @@ jobs: with: ref: ${{ inputs.ref }} fetch-depth: 0 + persist-credentials: false - name: Install system dependencies run: | @@ -137,12 +138,16 @@ jobs: - name: Run Benchmarks run: | + # Gate this PR job on directly measured runtimes only. uvx --with virtualenv asv continuous \ --launch-method spawn \ --interleave-rounds \ --append-samples \ --no-only-changed \ - -e -b Fast \ + -e \ + -b 'Fast\w*\.time_' \ + -b 'Fast(?!NewtonOverhead)\w*\.track_simulate(\(|$)' \ + -b 'Fast\w*\.track_(mean|p95)_\w*(time|ms)(\(|$)' \ ${{ inputs.base_ref }} \ ${{ inputs.ref }} continue-on-error: true diff --git a/.github/workflows/pr_auto_assign_creator.yml b/.github/workflows/pr_auto_assign_creator.yml index 91ef9ede00..41985927b0 100644 --- a/.github/workflows/pr_auto_assign_creator.yml +++ b/.github/workflows/pr_auto_assign_creator.yml @@ -49,11 +49,16 @@ jobs: } try { - await github.rest.issues.checkUserCanBeAssigned({ - owner, - repo, - assignee: author, - }); + // PR participants may be assignable without repository-wide access. + await github.request( + 'GET /repos/{owner}/{repo}/issues/{issue_number}/assignees/{assignee}', + { + owner, + repo, + issue_number: pr.number, + assignee: author, + } + ); } catch (error) { if (error.status === 404) { core.warning( diff --git a/.github/workflows/pr_target_aws_gpu_benchmarks.yml b/.github/workflows/pr_target_aws_gpu_benchmarks.yml index f10b57e0ef..a6e6902d19 100644 --- a/.github/workflows/pr_target_aws_gpu_benchmarks.yml +++ b/.github/workflows/pr_target_aws_gpu_benchmarks.yml @@ -96,7 +96,17 @@ jobs: needs: - check-author-membership - require-approval - if: github.repository == 'newton-physics/newton' && (!cancelled()) + if: >- + github.repository == 'newton-physics/newton' && + !cancelled() && + needs.check-author-membership.result == 'success' && + ( + needs.require-approval.result == 'success' || + ( + needs.require-approval.result == 'skipped' && + needs.check-author-membership.outputs.membership_status == 'CONFIRMED_MEMBER' + ) + ) uses: ./.github/workflows/aws_gpu_benchmarks.yml with: ref: ${{ github.event.pull_request.head.sha }} diff --git a/.github/workflows/pr_target_aws_gpu_tests.yml b/.github/workflows/pr_target_aws_gpu_tests.yml index 800c7f0d12..6380e86479 100644 --- a/.github/workflows/pr_target_aws_gpu_tests.yml +++ b/.github/workflows/pr_target_aws_gpu_tests.yml @@ -96,7 +96,17 @@ jobs: needs: - check-author-membership - require-approval - if: github.repository == 'newton-physics/newton' && (!cancelled()) + if: >- + github.repository == 'newton-physics/newton' && + !cancelled() && + needs.check-author-membership.result == 'success' && + ( + needs.require-approval.result == 'success' || + ( + needs.require-approval.result == 'skipped' && + needs.check-author-membership.outputs.membership_status == 'CONFIRMED_MEMBER' + ) + ) uses: ./.github/workflows/aws_gpu_tests.yml with: ref: ${{ github.event.pull_request.head.sha }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 6c91145bff..2737fe219a 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -15,7 +15,6 @@ jobs: runs-on: ubuntu-latest permissions: contents: read - actions: write steps: - name: Harden the runner (Audit all outbound calls) uses: step-security/harden-runner@fa2e9d605c4eeb9fcad4c99c224cee0c6c7f3594 # v2.16.0 diff --git a/AGENTS.md b/AGENTS.md index 9efa5de336..805a9a4c68 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -40,6 +40,7 @@ uv run --extra dev --extra torch-cu12 -m newton.tests # with Py ### Testing guidelines +- Give every test function or method a docstring using triple double quotes (`"""..."""`). Start with a concise one-line summary in imperative mood that states what the test verifies. For a particularly complex test, add a body that elaborates on the tested behavior, separated from the summary by a blank line following Google-style docstring conventions. - Never call `wp.synchronize()` or `wp.synchronize_device()` right before `.numpy()` on a Warp array. This is redundant as `.numpy()` performs a synchronous device-to-host copy that completes all outstanding work. ```bash diff --git a/CHANGELOG.md b/CHANGELOG.md index fa10071e81..c744b36c56 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,34 +4,65 @@ ### Added +- Break the viewer's shape count down into visual and collision shapes. The two are listed under `Shapes` in the stats overlay and need not sum to the total, since a shape can be both. +- Add selection of the shapes included in model shape BVHs through `Model.bvh_build_shapes(shape_flags=...)` and `ModelBuilder.default_bvh_cfg.shape_flags`, e.g. `ShapeFlags.VISIBLE | ShapeFlags.COLLIDE_SHAPES` to also include collision shapes. +- Add a `damping` parameter to `ModelBuilder.add_joint_ball()` that applies passive angular damping to all three ball-joint DOFs; when omitted, `ModelBuilder.default_joint_cfg.damping` applies. +- Add per-world `xforms` argument to `ModelBuilder.replicate()` for batching explicitly positioned worlds. - Import USD deformable bodies in `ModelBuilder.add_usd()` (experimental; based on the proposed [AOUSD Deformable Body Physics schema](https://github.com/aousd/OpenUSD-proposals/blob/5d89c0ed46a26de92f4d3fefef3bfad6500c07ce/proposals/physics_deformables/wp_deformable_physics.md)). Curves become cables (capsule bodies joined by cable joints, each cable wrapped in its own articulation), meshes become cloth (FEM triangles with bending edges), and tet meshes become soft bodies. Bound deformable materials supply thickness, stiffness, and density, and the proposal's mass precedence (per-point `physics:masses`, then body `mass`/`density`, then material density) is honored. Each imported deformable's element ranges and authored material attributes are returned by prim path when the experimental `add_usd(..., return_deformable_results=True)` option is passed; the returned maps are build-time snapshots (not live selections), and the default return shape carries no deformable entries. Collision participation follows `PhysicsCollisionAPI` / `physics:collisionEnabled`: cables without an enabled collider import without collision, while cloth and volume deformables warn that particle collision cannot be disabled yet. Standard `physics:filteredPairs` pairs involving a cable expand to its segment shapes; pairs naming a cloth or volume deformable warn and are not lowered. Disabled or kinematic deformables and malformed topology warn and are skipped. Cable and cloth material attributes are returned as authored in `path_cable_attrs` / `path_cloth_attrs`, so solvers with richer cable or cloth models can rebuild from the import. See the USD parsing documentation for the supported subset and limitations. - Import AOUSD proposal `PhysicsAttachment` prims for cables in `ModelBuilder.add_usd()`. Cable `point` / `segment` sites with an `xform` target become hard ball joints, returned in `path_attachment_map` (with `return_deformable_results=True`); cloth/volume attachment sites warn and are kept in `path_attachment_attrs`. A hard, coincident `point`->`point` attachment between two cables welds them into one rod graph; a springy or non-coincident junction warns and is kept as data instead of welded. - Import AOUSD proposal `PhysicsElementCollisionFilter` prims in `ModelBuilder.add_usd()`: collisions between the paired element groups of `src0` and `src1` are filtered (group counts pair element-wise; a count of `0` or an empty counts array selects all elements). Sources resolve to imported cables, rigid bodies, or collider prims; cloth/volume element sources warn and are skipped. - Add scalar value-based OpenCV, F-theta, and Kannala-Brandt fisheye camera ray helpers to `SensorTiledCamera.utils`, plus pinhole aperture/focal-length parameters, `compute_camera_transforms_usd()`, `compute_camera_rays_usd_pinhole()`, and optional preallocated ray output writes. +- Add CUDA-graph-capturable rebuildable sparse grids to `SolverImplicitMPM` when `max_active_cell_count` is positive, with optional `max_leaf_node_count`, `max_lower_node_count`, and `max_upper_node_count` hierarchy capacities. +- Add opt-in isolated multi-world implicit MPM with capacity-bounded rebuildable sparse grids, selective world resets, outer graph capture, and asynchronous overflow reporting; legacy shared topology remains the default. - Add `cloth_stiff_material_hanging` and `cloth_stiff_material_stretch` examples regression-guarding the new Neo-Hookean triangle material (stability under gravity at extreme stiffness, and bulk area-preservation across a Poisson-ratio sweep) +- Add contact examples for Newton's cradle, a balance bird, and a domino spiral - Add `ViewerUSD(points_as_spheres=...)` to render `log_points` particles as a `UsdGeom.PointInstancer` of sphere prototypes; enabled by default (opt out with `points_as_spheres=False` for flat `UsdGeom.Points` splats) - Add list-of-pattern and explicit-index selectors to `ArticulationView`. - Add `newton[onnx]` for ONNX policy inference through Warp-NN; `ControllerNeuralMLP`, `ControllerNeuralLSTM`, and RL policy examples can run exported `.onnx` policies without requiring PyTorch for ONNX execution. - Add three VBD contact examples — `vbd_rigid_rigid_contact`, `vbd_soft_rigid_contact`, and `vbd_soft_rigid_mix_contact` — demonstrating rigid-rigid, soft (particle-rigid), and mixed cloth-bag contacts - Add masked rigid-body reset support to `SolverVBD`; particle resets are not yet supported. (#3256) - Add viewer layer system to overlay multiple solvers/models in supported rendering viewers; call `ViewerBase.activate(layer_id)` to route subsequent `set_model` / `log_state` / `log_*` calls into a named layer, `ViewerBase.set_layer_visible()` to toggle layers independently, and `ViewerBase.set_layer_transform()` to position layers side-by-side. See `example_basic_multi_solver_overlay.py` +- Add `Heightfield.create_from_mesh()` and `newton.utils.rasterize_mesh_to_heightfield()` to build a heightfield collider by ray-casting a `wp.Mesh`, replacing a large static terrain mesh with an equivalent heightfield. - Add `ViewerBase.camera_speed` to configure keyboard translation speed for interactive viewers. (#3439) +- Add a "Show Ground" visualization toggle (`ViewerBase.show_ground`, default on) to hide or show ground-plane shapes in the viewer. +- Add opt-in DVI forward dynamics to `SolverKamino` through `SolverKamino.Config(dynamics_solver="dvi")`, with sparse and dense execution, DVI-specific diagnostics, and warm-starting. PADMM remains the default. +- Add opt-in DVI forward dynamics to `SolverKamino` through `SolverKamino.Config(dynamics_solver="dvi")`, with sparse and dense execution, DVI-specific convergence diagnostics, warm-starting, bounded contact-recovery controls, and RCM-reordered bilateral factorization with reusable ordering and panel-parallel numeric factorization for large systems. PADMM remains the default. - Add SDF contact support for convex-hull shapes with mesh-attached SDFs and opt-in SDF contact generation for box shapes. +- Warn in `ModelBuilder.add_usd()` when a rigid body prim has a mirrored (negative-determinant) world transform. Improper transforms have no unique rotation decomposition, so imported body and joint frames can acquire a spurious constant rotation (common with mirror-scaled CAD exports); the warning recommends baking the reflection into the mesh geometry before import. - Add opt-in filtering of static-static, static-kinematic, and kinematic-kinematic contacts during broad-phase collision detection. Set `CollisionPipeline(include_static_kinematic_pairs=False)` to enable filtering; the default preserves existing contact generation. `Model.shape_contact_pairs` remains an unfiltered superset for direct consumers such as `SolverKamino` and hydroelastic SDF setup. - Add opt-in `body_frame_origin="com"` to `ModelBuilder.add_rod()` and `ModelBuilder.add_rod_graph()` for COM-centered cable capsule body frames. - Add `sign_method` argument to `Mesh.build_sdf` and `SDF.create_from_mesh` support for a `"normal"` (angle-weighted pseudo-normal) sign strategy, for selecting the inside/outside sign of the baked SDF (`"auto"`, `"parity"`, `"winding"`, or `"normal"`). - Add `forward_depth_image` output support to `SensorTiledCamera.update()` and `SensorTiledCamera.utils.create_forward_depth_image_output()` for native forward-depth rendering without post-processing `depth_image`. - +- Add optional `shear_stiffness`/`shear_damping` and `twist_stiffness`/`twist_damping` controls to `ModelBuilder.add_joint_cable()`, `ModelBuilder.add_rod()`, and `ModelBuilder.add_rod_graph()`; omitted shear defaults to stretch and omitted twist defaults to bend for compatibility. +- Add `newton.utils.CableStiffness` and extend `newton.utils.create_cable_stiffness_from_elastic_moduli()` with `poissons_ratio`/`shear_modulus` inputs that include torsional `GJ/L` stiffness. +- Add VBD cable validation examples covering bend stiffness, analytical bend/twist response, torsion material mapping, routed twist transfer, twist-buckling link verification, Michell/Zajac threshold behavior, and Dahl hysteresis. +- Add a cable plectoneme example demonstrating twist-driven supercoiling with self-contact. +- Import authored USD cable stretch, shear, bend, and twist stiffness independently through `ModelBuilder.add_usd()`. +- Add compiled regular-expression support to label-based selectors while preserving glob strings. +- Add simulation throughput, real-time factor, p95 step-time, steady-state GPU-memory, timestep, and MuJoCo solver-iteration metrics to the ASV robot-learning benchmarks. +- Add `joint_dof_mask` to `newton.ik.IKSolver` to keep selected joint DOFs fixed during LM optimization. (#3488) ### Changed +- Decide collider visibility from USD `purpose` and visibility rather than from a bound render material. A collider whose `purpose` resolves to `default` is viewport geometry and is drawn; mark it `guide` to state that it is collision-only. Previously an unrelated visual elsewhere in the scene could make a collider vanish. `force_show_colliders` and `hide_collision_shapes` are unchanged. +- Disable the implicit positive Dahl-friction defaults in `SolverVBD.register_custom_attributes()` (deprecated in 1.3.0): `vbd:dahl_eps_max` and `vbd:dahl_tau` now default to zero, and Dahl cable friction is enabled only where both are authored positive. Pass `dahl_defaults_enabled=True` to temporarily restore the old defaults; the compatibility mode will be removed in a future release. +- Keep the authored render mesh when `ModelBuilder.add_usd()` approximates a collider. `physics:approximation` is scoped to collision, so a Mesh that is both render geometry and a collider now imports as an approximated collision shape plus a visual shape carrying the original topology, instead of replacing the render mesh with the approximation. This raises `Model.shape_count` for such prims: iterate on `ShapeFlags.COLLIDE_SHAPES` rather than assuming one shape per collider prim. The visual shape adds no mass and no collision, appends after the originals so existing shape indices and `path_shape_map` entries are unchanged, and is skipped when `load_visual_shapes=False`. - Compile tiled camera render kernels with CUDA fast math by default for faster rendering; set `SensorTiledCamera.render_config.enable_fast_math = False` for bit-exact, IEEE-precise output. +- Make `CollisionPipeline` the sole owner of rigid-contact geometry for `SolverVBD`: `"latest"` supplies fresh geometry and `"sticky"` supplies replayed geometry. `SolverVBD(rigid_contact_history=True)` uses either mode's match indices only to warm-start its numeric lambda/penalty state. - Optimize raycast/raytrace queries by restructuring ray-shape intersection into local-space primitives and compile specialized depth/shadow variants that skip unused surface-normal work (mesh shadows also use any-hit queries). +- Change experimental `SolverVBD` cable constraint slots from `[STRETCH=0, BEND=1]` to `[STRETCH=0, SHEAR=1, BEND=2, TWIST=3]`, allowing each stiffness and constraint mode to be configured independently. Existing cable calls using raw `slot=1` or `JointSlot.ANGULAR` now select shear; use `JointSlot.BEND` (now slot 2) to select bending. +- Map `shape_material_kf` to per-contact MuJoCo `solreffriction` in `SolverMuJoCo` (elliptic friction cones with Newton contacts); resolve `kf` with priority/`solmix`, treat a resolved `kf = 0` as frictionless, and use native MuJoCo contacts or a pyramidal cone to preserve the previous solref-inherited friction. +- Load visual-only USD geometry outside rigid-body hierarchies as static shapes by default; pass `load_static_visual_shapes=False` to retain the previous body-associated-visuals-only behavior. - Improve `SolverKamino` GPU simulation and kernel compilation performance. +- Speed up `Mesh.create_heightfield()` and `Mesh.create_terrain()` by building the vertex and index buffers in place, substantially reducing construction time and peak memory for large terrain grids such as those used by Isaac Lab. +- Load solver backends lazily on first access to speed up `import newton`; access solver classes through `newton.solvers` as before, and import solver modules explicitly if module-level side effects are required. +- Speed up USD mesh import for faceVarying normals by resolving the common single-cluster case for all vertices at once instead of clustering every face corner in Python; the split vertices, indices, normals, and UVs are unchanged, except that a corner sitting exactly at `vertex_splitting_angle_threshold_deg` from its cluster may now cluster differently. - Speed up `ModelBuilder.replicate()` for large world counts by merging all copies in one pass; it no longer calls `add_world()` or `add_builder()` per copy, so `ModelBuilder` subclass overrides of those methods are not invoked during replication. +- Treat `BodyFlags.KINEMATIC` bodies as zero-effective-mass implicit-MPM colliders when `SolverImplicitMPM.setup_collider()` is called without `body_mass`. Pass an explicit `body_mass` array to override the model-derived collider masses. ### Deprecated - Deprecate scalar `ModelBuilder.gravity`; pass a three-component gravity vector instead. +- Deprecate and ignore `SolverVBD`'s `rigid_contact_stick_motion_eps`, `rigid_contact_stick_freeze_translation_eps`, and `rigid_contact_stick_freeze_angular_eps`; use collision-pipeline sticky matching for persistent geometry. The SolverVBD body deadzone was removed without replacement. - Deprecate per-DOF `newton:{axis}:limitStiffness` and `newton:{axis}:limitDamping` attributes (where `{axis}` is `linear`, `angular`, `rotX`, `rotY`, or `rotZ`). Use the broadcast `newton:limitStiffness` and `newton:limitDamping` attributes from `NewtonJointAPI` instead; the broadcast value applies uniformly to all DOFs on the joint. For joints requiring per-DOF variance, split into separate 1-DOF (revolute / prismatic) joints. - Deprecate passing solver constructor options positionally after stable positional inputs such as `model` and explicit solver configs; migrate calls such as `SolverVBD(model, 10)` to `SolverVBD(model, iterations=10)`. - Deprecate `Model.contacts()` and `Model.collide()` in favor of explicitly creating a `CollisionPipeline`, allocating with `pipeline.contacts()`, and detecting collisions with `pipeline.collide(state, contacts)`. @@ -45,28 +76,62 @@ - Deprecate omitting `body_frame_origin` in `ModelBuilder.add_rod()` and `ModelBuilder.add_rod_graph()`; the implicit behavior still uses the existing start-node body-frame convention during the deprecation window, but the implicit default will change to `body_frame_origin="com"` in a future release. Pass `body_frame_origin="start"` to preserve the legacy frame or `body_frame_origin="com"` to opt into the future COM-centered frame. - Deprecate mutating `Model.shape_collision_filter_pairs`; modify `ModelBuilder.shape_collision_filter_pairs` before calling `finalize()` and rebuild the model instead, because mutating finalized collision filters does not rebuild `Model.shape_contact_pairs`. - Deprecate reading legacy vendor-namespaced deformable material attributes (`omniphysics:`, `physxDeformableBody:`) off any bound material in `newton.usd.get_tetmesh()`, `newton.TetMesh.create_from_usd()`, and `ModelBuilder.add_usd()`. They are still read during the deprecation window, with a `DeprecationWarning`; a future release will read only canonical `physics:` attributes from a material applying `PhysicsVolumeDeformableMaterialAPI`. Migrate by authoring the canonical attributes, or keep the old behavior without the warning via `compat_namespaces=newton.usd.DEFORMABLE_LEGACY_NAMESPACES` (`get_tetmesh` / `create_from_usd`) or `schema_resolvers=[..., SchemaResolverPhysx()]` (`add_usd`). `compat_namespaces` is now keyword-only; pass `()` to opt into the canonical-only behavior today. -- Deprecate the `indices` argument of `MeshAdjacency` in favor of `tri_indices` -- Deprecate `MeshAdjacency.add_edge`; construct a `MeshAdjacency` with `edge_indices` (`[o0, o1, v0, v1]` rows) instead - Deprecate implicit render-config updates in `SensorTiledCamera.utils.create_default_light()` and `SensorTiledCamera.utils.assign_checkerboard_material()`; set `sensor.default_render_config.enable_shadows` or `sensor.default_render_config.enable_textures` explicitly instead. - Deprecate `SensorTiledCamera(..., config=...)` in favor of `SensorTiledCamera(..., default_render_config=...)`; migrate constructor calls that pass a render config to the new keyword. - Deprecate `SensorTiledCamera.render_config` in favor of `SensorTiledCamera.default_render_config`; migrate `sensor.render_config.enable_shadows = True` to `sensor.default_render_config.enable_shadows = True`. - Deprecate `SensorTiledCamera.utils.compute_pinhole_camera_rays()` in favor of `SensorTiledCamera.utils.compute_camera_rays_pinhole()`. +### Removed + +- Remove the deprecated SDF compatibility attributes `Model.shape_sdf_index`, `Model.texture_sdf_data`, `Model.texture_sdf_coarse_textures`, `Model.texture_sdf_subgrid_textures`, `Model.texture_sdf_subgrid_start_slots`, `Model.sdf_block_coords`, `Model.sdf_index2blocks`, and `SDF.texture_block_coords` (deprecated in 1.3.0); the hydroelastic broadphase derives block coordinates arithmetically and the remaining storage is internal. +- Remove the deprecated `newton.geometry.build_bvh_shape()`, `refit_bvh_shape()`, `build_bvh_particle()`, and `refit_bvh_particle()` helpers (deprecated in 1.3.0); use `Model.bvh_build_shapes()`, `Model.bvh_refit_shapes()`, `Model.bvh_build_particles()`, and `Model.bvh_refit_particles()` instead. +- Remove the deprecated `Model.has_heightfields` property (deprecated in 1.3.0); use `Model.heightfield_count`, or `model.heightfield_count > 0` for boolean checks, instead. +- Remove the deprecated `SolverNotifyFlags` enum (deprecated in 1.3.0); use `ModelFlags` instead. +- Remove the deprecated `ls_parallel` parameter of `SolverMuJoCo` (deprecated in 1.3.0); parallel line search was removed from `mujoco_warp` and the option had no effect. + ### Fixed +- Convert `newton:mimicCoef0` from degrees to radians when the mimic follower joint is angular. Assets authored against the old behavior need the value rescaled to degrees. +- Complete Kamino RCM traversal for large and disconnected systems and reuse the resulting permutation by default; set `reuse_permutation=False` to recompute it for changing matrix topology. +- Fix panel-parallel RCM-blocked LLT factorization hanging when a matrix ends in a partial tile. +- Fix `ModelBuilder.add_usd()` marking a `guide`-purpose collider visible when it has a bound render material. Such a collider is not viewport geometry, and the extra `VISIBLE` flag left it drawn by the viewer's visual toggle instead of its collision toggle. `force_show_colliders` still reveals it. +- Fix USD capsule, cylinder, and cone visual and site scaling to follow the authored primitive axis. +- Fix USD plane visual width and length to scale along the axes defined by the `UsdGeomPlane` schema, and orient X- and Y-axis plane visuals along the authored axis. +- Validate `ArticulationView` mask shapes and devices before launching selection kernels. (#3448) +- Exclude active particles with non-finite positions from rebuildable `SolverImplicitMPM` sparse-grid packing. +- Fix masked `SolverCoupledProxy.reset()` calls clearing proxy feedback history for unselected worlds. +- Fix hydroelastic primitive texture SDF generation to sample analytic primitive distances instead of temporary tessellated meshes. (#3239) - Fix MJCF, URDF, and USD imports rendering collision-only bodies as visuals when the asset authors visual geometry elsewhere. (#3291) - Fix `SchemaResolverPhysx` reading every D6 translational limit gain from the `linear` instance instead of its `transX`, `transY`, or `transZ` instance. +- Fix USD capsule, cylinder, and cone visuals and sites without authored `radius`/`height` to use the UsdGeom schema fallbacks, matching collision shapes. - Fix `ViewerUSD` texture consumers observing partially written PNGs by publishing generated textures atomically (#3288) +- Fix loading of textures packaged inside `.usdz` archives; package-relative asset paths such as `scene.usdz[tex.png]` are resolved through USD's asset resolver instead of being treated as filesystem paths. +- Fix `ModelBuilder.add_usd()` raising `ValueError` when importing a mesh whose material subset binds a texture that decodes to an image array. +- Fix textured USD visual meshes rendering tinted by the default per-shape palette color; a textured mesh without an authored scalar color now imports with a white base color so its texture is shown untinted. +- Fix `ModelBuilder.add_usd()` selecting a non-color map (e.g. a roughness, metallic, or normal map) as a mesh's base-color texture. A connected `UsdUVTexture` is now accepted only when it feeds a base-color input by name through its multi-channel color output, and a shader's direct-asset color parameter (e.g. an MDL `diffuse_texture`) is likewise identified by name. +- Fix scrambled textures on USD meshes whose texture-coordinate primvar is not named `st` (e.g. `st_0`). The texcoord set is now resolved from the bound material's shader network (the `UsdPreviewSurface` texture reader's `varname` or an MDL/OmniPBR `uv_space_index`), and textured material subsets slice real per-corner UVs and authored normals instead of collapsing faceVarying data per vertex. - Fix builder merging (`ModelBuilder.add_builder()`, `add_world()`, `replicate()`) offsetting negative reference sentinels in custom attribute values stored as NumPy or Warp integer scalars. - Fix `ModelBuilder.add_usd()` requiring the optional `mujoco` package when handling `MjcActuator` prims, including during default MJC equality conversion. +- Fix `ModelBuilder.add_usd()` ignoring enabled collider mass properties and counting disabled colliders toward body mass. (#3594) - Report malformed MJCF free-joint and inertial inputs with deterministic validation errors, and ignore MJCF mesh geom `size` lengths consistently. +- Fix MJCF imports ignoring material and inline RGBA colors on primitive geoms. +- Fix `SolverVBD` failing to construct on large multi-world scenes containing particles and rigid shapes. (#3660) - Fix Style3D solver divergence caused by isolated vertices. +- Fix compiler warnings about overflowing int32 constants when compiling SDF texture and `SensorTiledCamera` kernels. +- Fix USD site import to discover sites beneath non-visual containers, collider prims, and instanceable rigid-body prims independently of `load_visual_shapes`; the reworked traversal also speeds up import of scenes with many nested `Xform` or instance prims. +- Fix `SolverFeatherstone` BALL joints to apply passive `joint_damping` on all three angular DOFs. +- Fix `eval_ik()` and `SolverSemiImplicit` rounding small float32 revolute-joint angles to zero. (#3434) - Fix excessive memory usage when importing MJCF or URDF models containing many visual-only shapes with self-collisions disabled. +- Fix `FastKitchenG1` ASV metrics to build the kitchen scene instead of a plain G1 model. - Fix the `diffsim_bear` example crashing with its default CUDA configuration and diverging after a few training iterations. - Fix masked PID state reset to execute on the integral-state device. (#3447) +- Reject invalid hollow primitive shell thickness before computing inertia. +- Fix `ModelBuilder.add_mjcf()` ignoring positive explicit mass on mesh geoms. (#3595) - Preserve muscles and rigid-body color groups when copying or replicating a `ModelBuilder`. - Fix `ModelBuilder.add_usd()` to honor `PhysicsScene.gravityDirection`, including stage-to-builder rotation and per-world imports. +- Fix `ModelBuilder.add_mjcf()` to honor compiler `inertiafromgeom` and `inertiagrouprange`, and keep inferred mass independent of `parse_visuals`. (#3596) - Fix stale overlay layers remaining visible after switching examples in the OpenGL viewer. +- Fix `SolverKamino` CG/CR solves silently under-iterating on CPU graph capture; the capture-safe loop path now runs on any capturing device, not only CUDA, so CPU captures no longer record a stale host-readback convergence decision at record time. - Reject incompatible custom attribute and frequency definitions before composing `ModelBuilder` instances. - Fix `cloth_franka` example rendering particles at simulation scale (cm) instead of viewer scale (m) - Fix `ModelBuilder` merges to accept array-valued transform fields and plain-list particle color groups. @@ -170,6 +235,7 @@ - Deprecate reading legacy vendor-namespaced deformable material attributes (`omniphysics:`, `physxDeformableBody:`) off any bound material in `newton.usd.get_tetmesh()`, `newton.TetMesh.create_from_usd()`, and `ModelBuilder.add_usd()`. They are still read during the deprecation window, with a `DeprecationWarning`; a future release will read only canonical `physics:` attributes from a material applying `PhysicsVolumeDeformableMaterialAPI`. Migrate by authoring the canonical attributes, or keep the old behavior without the warning via `compat_namespaces=newton.usd.DEFORMABLE_LEGACY_NAMESPACES` (`get_tetmesh` / `create_from_usd`) or `schema_resolvers=[..., SchemaResolverPhysx()]` (`add_usd`). `compat_namespaces` is now keyword-only; pass `()` to opt into the canonical-only behavior today. (#3192) - Deprecate the `indices` argument of `MeshAdjacency` in favor of `tri_indices`. (#3194) - Deprecate `MeshAdjacency.add_edge`; construct a `MeshAdjacency` with `edge_indices` (`[o0, o1, v0, v1]` rows) instead. (#3194) +- Deprecate the `MeshAdjacency.edges` dict accessor; use the `edge_indices` / `edge_tri_indices` arrays instead. (#3194) - Deprecate `SensorTiledCamera.utils.compute_pinhole_camera_rays()` in favor of `SensorTiledCamera.utils.compute_camera_rays_pinhole()`. (#3026) ### Fixed @@ -179,6 +245,10 @@ - Tune VBD contact settings in the `basic_shapes` and `cable_bundle_hysteresis` examples for more consistent friction and recovery behavior. (#3446) - Fix USD import ignoring ancestor material bindings with `strongerThanDescendants` strength when a mesh authors `material:binding` without applying `MaterialBindingAPI`: material resolution now uses UsdShade's canonical `ComputeBoundMaterial` unconditionally, which also adds collection-based binding support. Prims authoring bindings without the applied schema are invalid USD and now surface USD's own warning (once per prim per import) — fix such assets with `usdchecker` or `usd-validation-nvidia`. (#3350) - Fix `ModelBuilder.add_usd()` to honor `ignore_paths` in the custom-frequency traversal, so prims under ignored subtrees no longer register spurious custom-frequency rows in two-pass import workflows. (#3406) +- Reject inconsistent per-particle array lengths during bulk model construction and finalization. (#3458) +- Fix USD import topology depending on material vocabulary: mesh subsets now split on the authored material-binding structure, so a subset bound to a material whose properties Newton does not recognize (e.g. an MDL shader) imports as its own unshaded submesh instead of changing the mesh's imported shape count. +- Fix USD joint `physics:collisionEnabled` import so joints with two explicit bodies honor authored collision behavior; joints to world continue to allow body/world collisions, and articulation-wide self-collision filtering remains additive. +- Fix `ViewerFile.is_running()` to return `False` after `ViewerFile.close()` so headless recording loops can terminate like interactive viewers. (#3094) - Fix USD joint `physics:collisionEnabled` import so joints with two explicit bodies honor authored collision behavior; joints to world continue to allow body/world collisions, and articulation-wide self-collision filtering remains additive. (#3320) - Fix `ViewerFile.is_running()` to return `False` after `ViewerFile.close()` so headless recording loops can terminate like interactive viewers. (#3190; fixes #3094) - Fix mesh-approximation fallback behavior: @@ -208,7 +278,9 @@ - Fix `SolverMuJoCo` inertia randomization for bodies initialized with diagonal inertia. - Fix `SolverMuJoCo` static worldbody geometry poses so offset batched worlds collide with their local geometry. - Fix memory growth in the Style3D solver when CUDA Graph capture is disabled +- Fix `SolverMuJoCo` site poses for offset batched worlds and site poses and sizes for runtime shape updates. - Limit mouse-picking torque using each body's rotational inertia to prevent unstable angular acceleration on low-inertia bodies. +- Reject invalid `ModelBuilder.ShapeConfig` SDF and density values during shape validation. - Fix multi-angular-DOF `JointType.D6` kinematics and dynamics: (#2975) - Build `newton.eval_jacobian`, `SolverFeatherstone`, and IK analytic Jacobian angular motion-subspace columns in the current joint frame, so `J @ joint_qd` matches `State.body_qd` at non-identity configurations. - Report the correct `SolverMuJoCo` `State.body_qd` angular velocity at non-identity configurations. @@ -496,7 +568,6 @@ - Deprecate the top-level `Model.equality_constraint_*` arrays and `Model.equality_constraint_count`, the `ModelBuilder.equality_constraint_*` accumulators, `ModelBuilder.add_equality_constraint{,_connect,_weld,_joint}()`, and the `Model.AttributeFrequency.EQUALITY_CONSTRAINT` enum, in favor of the namespaced `model.mujoco.equality_constraint_*` fields (custom attributes on the `"mujoco:equality_constraint"` frequency). Migrate reads and writes to `model.mujoco.equality_constraint_*`, and construct rows via `ModelBuilder.add_custom_values(**{"mujoco:equality_constraint_*": ...})`. The deprecated names forward to the namespace during the deprecation window and will be removed in a future release. - Deprecate `SensorRaycast` in favor of `SensorTiledCamera`; migrate to `SensorTiledCamera.utils.compute_camera_rays_pinhole()` and `create_depth_image_output()` for single-camera depth rendering — see the `SensorRaycast` class docstring for a complete migration example - Deprecate and ignore `rigid_enable_dahl_friction` in `SolverVBD`; Dahl friction is now auto-detected from model attributes (`model.vbd.dahl_eps_max` / `model.vbd.dahl_tau`) -- Deprecate the `MeshAdjacency.edges` dict accessor; use the `edge_indices` / `edge_tri_indices` arrays instead - Deprecate `newton-actuators` package dependency; all actuator functionality is now built into `newton.actuators`. The dependency is kept for backward compatibility and will be removed in a future release; migrate imports from `newton_actuators` to `newton.actuators` ### Fixed diff --git a/README.md b/README.md index 18f5481951..7a37e83deb 100644 --- a/README.md +++ b/README.md @@ -283,6 +283,9 @@ If you run the examples from a source checkout with uv, use + + Cable Plectoneme + @@ -293,6 +296,7 @@ If you run the examples from a source checkout with uv, use python -m newton.examples cable_cross_slide_table + python -m newton.examples cable_plectoneme @@ -749,6 +753,9 @@ If you run the examples from a source checkout with uv, use + + Newton Cradle + @@ -758,6 +765,31 @@ If you run the examples from a source checkout with uv, use python -m newton.examples contacts_rj45_plug + + python -m newton.examples newton_cradle + + + + + + Balance Bird + + + + + Domino Spiral + + + + + + + + python -m newton.examples balance_bird + + + python -m newton.examples domino_spiral + diff --git a/asv.conf.json b/asv.conf.json index b82c81da56..1d5b490e62 100644 --- a/asv.conf.json +++ b/asv.conf.json @@ -16,7 +16,7 @@ "build_command": ["python -m build --wheel -o {build_cache_dir} {build_dir}"], "install_command": [ "in-dir={env_dir} python -m pip install {wheel_file}[dev] --extra-index-url=https://pypi.nvidia.com/", - "in-dir={env_dir} python -m pip install {wheel_file}[sim] warp-lang==1.16.0.dev20260716 mujoco==3.10.0 mujoco-warp==3.10.0.2 --extra-index-url=https://pypi.nvidia.com/", + "in-dir={env_dir} python -m pip install {wheel_file}[sim] warp-lang==1.16.0.dev20260716 mujoco==3.10.0 mujoco-warp==3.10.0.3 --extra-index-url=https://pypi.nvidia.com/", "python -m pip install torch==2.10.0+cu130 --index-url https://download.pytorch.org/whl/cu130" ] } diff --git a/asv/benchmarks/benchmark_kamino.py b/asv/benchmarks/benchmark_kamino.py index e9e69b0251..9096de8118 100644 --- a/asv/benchmarks/benchmark_kamino.py +++ b/asv/benchmarks/benchmark_kamino.py @@ -12,6 +12,11 @@ import newton +if __package__: + from .benchmark_metrics import validate_simulation_state +else: + from benchmark_metrics import validate_simulation_state + _NUM_ACTIONS = 12 _OBS_DIM = 94 _MIN_STANDING_HEIGHT = 0.20 @@ -435,37 +440,23 @@ def step(self): self.sim_time += self.frame_dt def test_final(self): - state_values = {} - for name in ("joint_q", "body_q", "body_qd"): - values = getattr(self.state_0, name).numpy() - if not np.isfinite(values).all(): - raise RuntimeError(f"Simulation produced non-finite values in state.{name}") - state_values[name] = values - - body_count = self.model.body_count // self.world_count - body_qd = state_values["body_qd"].reshape(self.world_count, body_count, 6) - max_linear_speed = np.linalg.norm(body_qd[:, :, :3], axis=-1).max() - max_angular_speed = np.linalg.norm(body_qd[:, :, 3:], axis=-1).max() - if max_linear_speed > _MAX_BODY_LINEAR_SPEED: - raise RuntimeError( - f"Maximum body linear speed is {max_linear_speed:.3f} m/s, exceeding {_MAX_BODY_LINEAR_SPEED:.1f} m/s" - ) - if max_angular_speed > _MAX_BODY_ANGULAR_SPEED: - raise RuntimeError( - f"Maximum body angular speed is {max_angular_speed:.3f} rad/s, " - f"exceeding {_MAX_BODY_ANGULAR_SPEED:.1f} rad/s" - ) + validate_simulation_state( + self.state_0, + max_linear_speed=_MAX_BODY_LINEAR_SPEED, + max_angular_speed=_MAX_BODY_ANGULAR_SPEED, + ) if self.policy_controller is None: return + body_count = self.model.body_count // self.world_count body_labels = [label.rsplit("/", 1)[-1] for label in self.model.body_label[:body_count]] try: pelvis_index = body_labels.index("pelvis") except ValueError as e: raise RuntimeError("DR Legs model has no pelvis root body") from e - body_q = state_values["body_q"].reshape(self.world_count, body_count, 7)[:, pelvis_index] + body_q = self.state_0.body_q.numpy().reshape(self.world_count, body_count, 7)[:, pelvis_index] body_com = self.model.body_com.numpy().reshape(self.world_count, body_count, 3)[:, pelvis_index] quat_vector = body_q[:, 3:6] twice_cross = 2.0 * np.cross(quat_vector, body_com) diff --git a/asv/benchmarks/benchmark_metrics.py b/asv/benchmarks/benchmark_metrics.py new file mode 100644 index 0000000000..6984dffce8 --- /dev/null +++ b/asv/benchmarks/benchmark_metrics.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +import math +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np +import warp as wp +from asv_runner.benchmarks.mark import skip_benchmark_if + + +@dataclass(frozen=True) +class SimulationMetrics: + """Metrics collected from one simulation benchmark configuration.""" + + mean_world_step_time_ms: float + world_steps_per_second: float + real_time_factor: float + p95_frame_time_ms: float + gpu_memory_mib: float + sim_dt: float + sim_substeps: int + solver_niter_mean: float | None = None + solver_niter_max: float | None = None + + +class _SimulationMetricTracks: + """ASV track methods backed by cached simulation metrics.""" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_simulate(self, metrics, world_count): + return metrics[world_count].mean_world_step_time_ms + + track_simulate.unit = "ms/world-step" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_simulation_steps_per_second(self, metrics, world_count): + return metrics[world_count].world_steps_per_second + + track_simulation_steps_per_second.unit = "world-steps/s" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_real_time_factor(self, metrics, world_count): + return metrics[world_count].real_time_factor + + track_real_time_factor.unit = "x" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_p95_step_time(self, metrics, world_count): + return metrics[world_count].p95_frame_time_ms + + track_p95_step_time.unit = "ms/frame" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_steady_state_gpu_memory(self, metrics, world_count): + return metrics[world_count].gpu_memory_mib + + track_steady_state_gpu_memory.unit = "MiB" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_sim_dt(self, metrics, world_count): + return metrics[world_count].sim_dt + + track_sim_dt.unit = "s" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_sim_substeps(self, metrics, world_count): + return metrics[world_count].sim_substeps + + track_sim_substeps.unit = "simulation-steps/frame" + + +class _SimulationMetricTracksUnparameterized: + """ASV track methods backed by one cached simulation configuration.""" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_mean_world_step_time(self, metrics): + return metrics.mean_world_step_time_ms + + track_mean_world_step_time.unit = "ms/world-step" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_simulation_steps_per_second(self, metrics): + return metrics.world_steps_per_second + + track_simulation_steps_per_second.unit = "world-steps/s" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_real_time_factor(self, metrics): + return metrics.real_time_factor + + track_real_time_factor.unit = "x" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_p95_step_time(self, metrics): + return metrics.p95_frame_time_ms + + track_p95_step_time.unit = "ms/frame" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_steady_state_gpu_memory(self, metrics): + return metrics.gpu_memory_mib + + track_steady_state_gpu_memory.unit = "MiB" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_sim_dt(self, metrics): + return metrics.sim_dt + + track_sim_dt.unit = "s" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_sim_substeps(self, metrics): + return metrics.sim_substeps + + track_sim_substeps.unit = "simulation-steps/frame" + + +def compute_simulation_metrics( + frame_times: Sequence[float], + sim_dt: float, + sim_substeps: int, + world_count: int, + gpu_memory_bytes: int, + experience_frame_times: Sequence[float] | None = None, +) -> SimulationMetrics: + """Compute comparable simulation metrics from synchronized frame times.""" + if not frame_times or any(not math.isfinite(value) or value <= 0.0 for value in frame_times): + raise ValueError("frame_times must contain positive finite values") + if experience_frame_times is None: + experience_frame_times = frame_times + if len(experience_frame_times) != len(frame_times) or any( + not math.isfinite(value) or value <= 0.0 for value in experience_frame_times + ): + raise ValueError("experience_frame_times must contain one positive finite value per frame") + if not math.isfinite(sim_dt) or sim_dt <= 0.0: + raise ValueError("sim_dt must be positive and finite") + if sim_substeps <= 0 or world_count <= 0: + raise ValueError("sim_substeps and world_count must be positive") + if gpu_memory_bytes < 0: + raise ValueError("gpu_memory_bytes must be non-negative") + + total_time = sum(frame_times) + experience_total_time = sum(experience_frame_times) + world_steps = len(frame_times) * sim_substeps * world_count + return SimulationMetrics( + mean_world_step_time_ms=total_time * 1000.0 / world_steps, + world_steps_per_second=world_steps / experience_total_time, + real_time_factor=world_steps * sim_dt / experience_total_time, + p95_frame_time_ms=float(np.percentile(experience_frame_times, 95.0)) * 1000.0, + gpu_memory_mib=gpu_memory_bytes / 1024**2, + sim_dt=sim_dt, + sim_substeps=sim_substeps, + ) + + +def validate_simulation_state( + state: Any, + max_linear_speed: float, + max_angular_speed: float, + quaternion_tolerance: float = 1.0e-3, +): + """Validate finite rigid-body state, normalized rotations, and bounded speeds.""" + state_values = {} + for name in ("joint_q", "joint_qd", "body_q", "body_qd"): + values = getattr(state, name).numpy() + if not np.isfinite(values).all(): + raise RuntimeError(f"Simulation produced non-finite values in state.{name}") + state_values[name] = values + + body_q = state_values["body_q"].reshape(-1, 7) + quaternion_norms = np.linalg.norm(body_q[:, 3:7], axis=-1) + if not np.allclose(quaternion_norms, 1.0, atol=quaternion_tolerance, rtol=0.0): + max_error = np.abs(quaternion_norms - 1.0).max() + raise RuntimeError(f"Maximum body quaternion norm error is {max_error:.3g}") + + body_qd = state_values["body_qd"].reshape(-1, 6) + max_measured_linear_speed = np.linalg.norm(body_qd[:, :3], axis=-1).max() + max_measured_angular_speed = np.linalg.norm(body_qd[:, 3:], axis=-1).max() + if max_measured_linear_speed > max_linear_speed: + raise RuntimeError( + f"Maximum body linear speed is {max_measured_linear_speed:.3f} m/s, exceeding {max_linear_speed:.1f} m/s" + ) + if max_measured_angular_speed > max_angular_speed: + raise RuntimeError( + f"Maximum body angular speed is {max_measured_angular_speed:.3f} rad/s, " + f"exceeding {max_angular_speed:.1f} rad/s" + ) + + +def collect_simulation_metrics( + create_workload: Callable[[], Any], + world_count: int, + num_frames: int, + samples: int, + synchronize: Callable[[], None] | None = None, + validate: Callable[[Any], None] | None = None, + timer: Callable[[], float] = time.perf_counter, +) -> SimulationMetrics: + """Collect simulation metrics using internal or synchronized wall timing.""" + frame_times = [] + experience_frame_times = [] + gpu_memory_bytes = None + sim_dt = None + sim_substeps = None + + wp.synchronize_device() + device = wp.get_device() + free_memory_before = device.free_memory + + for sample_index in range(samples): + workload = create_workload() + if sim_dt is None: + sim_dt = workload.sim_dt + sim_substeps = workload.sim_substeps + elif workload.sim_dt != sim_dt or workload.sim_substeps != sim_substeps: + raise ValueError("simulation parameters changed between samples") + + if synchronize is not None: + synchronize() + for _ in range(num_frames): + experience_start_time = timer() + benchmark_start_time = workload.benchmark_time if synchronize is None else None + workload.step() + if synchronize is not None: + synchronize() + experience_frame_time = timer() - experience_start_time + experience_frame_times.append(experience_frame_time) + frame_times.append( + experience_frame_time + if benchmark_start_time is None + else workload.benchmark_time - benchmark_start_time + ) + + if sample_index == 0: + wp.synchronize_device() + gpu_memory_bytes = free_memory_before - device.free_memory + if gpu_memory_bytes < 0: + raise RuntimeError("GPU free memory increased after workload initialization") + if validate is not None: + validate(workload) + + return compute_simulation_metrics( + frame_times=frame_times, + sim_dt=sim_dt, + sim_substeps=sim_substeps, + world_count=world_count, + gpu_memory_bytes=gpu_memory_bytes, + experience_frame_times=experience_frame_times, + ) diff --git a/asv/benchmarks/benchmark_mujoco.py b/asv/benchmarks/benchmark_mujoco.py index 19499d19d0..f39feb28c5 100644 --- a/asv/benchmarks/benchmark_mujoco.py +++ b/asv/benchmarks/benchmark_mujoco.py @@ -23,7 +23,14 @@ from newton.sensors import SensorContact from newton.utils import EventTracer +if __package__: + from .benchmark_metrics import validate_simulation_state +else: + from benchmark_metrics import validate_simulation_state + _NEW_LAYOUT_AVAILABLE = hasattr(newton, "use_coord_layout_targets") +_MAX_BODY_LINEAR_SPEED = 100.0 +_MAX_BODY_ANGULAR_SPEED = 500.0 def _target_q(owner): @@ -368,14 +375,17 @@ def __init__( nconmax=None, builder=None, cone=None, + fps=600, + sim_substeps=10, ): if _NEW_LAYOUT_AVAILABLE: newton.use_coord_layout_targets = True - fps = 600 + if fps <= 0 or sim_substeps <= 0: + raise ValueError("fps and sim_substeps must be positive") self.sim_time = 0.0 self.benchmark_time = 0.0 self.frame_dt = 1.0 / fps - self.sim_substeps = 10 + self.sim_substeps = sim_substeps self.contacts = None self.sim_dt = self.frame_dt / self.sim_substeps self.world_count = world_count @@ -482,17 +492,24 @@ def step(self): self.apply_waypoint_control() wp.synchronize_device() - start_time = time.time() - if self.use_cuda_graph: + start_time = time.perf_counter() + if self.use_cuda_graph and self.graph is not None: wp.capture_launch(self.graph) else: self.simulate() wp.synchronize_device() - end_time = time.time() + end_time = time.perf_counter() self.benchmark_time += end_time - start_time self.sim_time += self.frame_dt + def test_final(self): + validate_simulation_state( + self.state_0, + max_linear_speed=_MAX_BODY_LINEAR_SPEED, + max_angular_speed=_MAX_BODY_ANGULAR_SPEED, + ) + def render(self): if self.renderer is None: return @@ -562,8 +579,6 @@ def create_solver( nconmax=None, cone=None, ): - solver_iteration = solver_iteration if solver_iteration is not None else 100 - ls_iteration = ls_iteration if ls_iteration is not None else 50 solver = solver if solver is not None else ROBOT_CONFIGS[robot]["solver"] integrator = integrator if integrator is not None else ROBOT_CONFIGS[robot]["integrator"] njmax = njmax if njmax is not None else ROBOT_CONFIGS[robot]["njmax"] diff --git a/asv/benchmarks/simulation/bench_anymal.py b/asv/benchmarks/simulation/bench_anymal.py index 545441b474..6fd9613421 100644 --- a/asv/benchmarks/simulation/bench_anymal.py +++ b/asv/benchmarks/simulation/bench_anymal.py @@ -1,15 +1,54 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers # SPDX-License-Identifier: Apache-2.0 +import os +import sys + import warp as wp from asv_runner.benchmarks.mark import skip_benchmark_if wp.config.enable_backward = False wp.config.log_level = wp.LOG_WARNING -import newton -import newton.examples -from newton.examples.robot.example_robot_anymal_c_walk import Example +parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.append(parent_dir) + +from benchmark_metrics import ( + _SimulationMetricTracksUnparameterized, + collect_simulation_metrics, + validate_simulation_state, +) + +_NUM_FRAMES = 50 +_MIN_BASE_HEIGHT = 0.4 +_MIN_FORWARD_PROGRESS = 0.25 + + +def _create_example(num_frames): + import newton # noqa: PLC0415 + import newton.examples # noqa: PLC0415 + from newton.examples.robot.example_robot_anymal_c_walk import Example # noqa: PLC0415 + + if hasattr(newton.examples, "default_args"): + args = newton.examples.default_args() + else: + args = None + return Example(newton.viewer.ViewerNull(num_frames=num_frames), args) + + +def _validate_workload(workload): + validate_simulation_state( + workload.state_0, + max_linear_speed=10.0, + max_angular_speed=50.0, + ) + root_position = workload.state_0.joint_q.numpy()[:3] + if root_position[2] < _MIN_BASE_HEIGHT: + raise RuntimeError(f"ANYmal base height is too low after {_NUM_FRAMES} frames: {root_position[2]:.3f} m") + if root_position[1] < _MIN_FORWARD_PROGRESS: + raise RuntimeError( + f"ANYmal made insufficient forward progress after {_NUM_FRAMES} frames: {root_position[1]:.3f} m" + ) class FastExampleAnymalPretrained: @@ -17,12 +56,8 @@ class FastExampleAnymalPretrained: number = 1 def setup(self): - self.num_frames = 50 - if hasattr(newton.examples, "default_args"): - args = newton.examples.default_args() - else: - args = None - self.example = Example(newton.viewer.ViewerNull(num_frames=self.num_frames), args) + self.num_frames = _NUM_FRAMES + self.example = _create_example(self.num_frames) @skip_benchmark_if(wp.get_cuda_device_count() == 0) def time_simulate(self): @@ -31,6 +66,24 @@ def time_simulate(self): wp.synchronize_device() +class FastMetricsExampleAnymalPretrained(_SimulationMetricTracksUnparameterized): + num_frames = _NUM_FRAMES + samples = 3 + world_count = 1 + + def setup_cache(self): + if wp.get_cuda_device_count() == 0: + return None + return collect_simulation_metrics( + create_workload=lambda: _create_example(self.num_frames), + world_count=self.world_count, + num_frames=self.num_frames, + samples=self.samples, + synchronize=wp.synchronize_device, + validate=_validate_workload, + ) + + if __name__ == "__main__": import argparse @@ -38,6 +91,7 @@ def time_simulate(self): benchmark_list = { "FastExampleAnymalPretrained": FastExampleAnymalPretrained, + "FastMetricsExampleAnymalPretrained": FastMetricsExampleAnymalPretrained, } parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) diff --git a/asv/benchmarks/simulation/bench_implicit_mpm.py b/asv/benchmarks/simulation/bench_implicit_mpm.py new file mode 100644 index 0000000000..fdab7c3bf9 --- /dev/null +++ b/asv/benchmarks/simulation/bench_implicit_mpm.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +import warp as wp +from asv_runner.benchmarks.mark import SkipNotImplemented, skip_benchmark_if + +wp.config.enable_backward = False +wp.config.log_level = wp.LOG_WARNING + +import newton +from newton.solvers import SolverImplicitMPM + + +class ImplicitMPMSingleWorld: + """Track the fixed-grid single-world fast path independently of batching.""" + + number = 1 + repeat = 5 + rounds = 2 + + def setup(self): + device = wp.get_device() + if not device.is_cuda: + raise SkipNotImplemented + + builder = newton.ModelBuilder(gravity=0.0) + SolverImplicitMPM.register_custom_attributes(builder) + builder.add_particle_grid( + pos=wp.vec3(0.0), + rot=wp.quat_identity(), + vel=wp.vec3(0.0), + dim_x=14, + dim_y=14, + dim_z=14, + cell_x=0.025, + cell_y=0.025, + cell_z=0.025, + mass=0.01, + jitter=0.0, + radius_mean=0.0125, + custom_attributes={"mpm:young_modulus": 1.0e4, "mpm:poisson_ratio": 0.2}, + ) + self.model = builder.finalize(device=device) + + config = SolverImplicitMPM.Config() + config.grid_type = "fixed" + config.grid_padding = 3 + config.voxel_size = 0.05 + config.transfer_scheme = "pic" + config.integration_scheme = "pic" + config.solver = "jacobi" + config.max_iterations = 10 + config.tolerance = 0.0 + config.warmstart_mode = "none" + + self.solver = SolverImplicitMPM(self.model, config=config, enable_timers=False) + self.state_0 = self.model.state() + self.state_1 = self.model.state() + self.dt = 0.001 + + # Compile and populate persistent fixed-grid data before timing. Two + # steps preserve the state-buffer orientation used by each ASV repeat. + self.solver.step(self.state_0, self.state_1, None, None, self.dt) + self.solver.step(self.state_1, self.state_0, None, None, self.dt) + wp.synchronize_device(device) + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_step(self): + for _ in range(10): + self.solver.step(self.state_0, self.state_1, None, None, self.dt) + self.state_0, self.state_1 = self.state_1, self.state_0 + wp.synchronize_device() + + +if __name__ == "__main__": + from newton.utils import run_benchmark + + run_benchmark(ImplicitMPMSingleWorld) diff --git a/asv/benchmarks/simulation/bench_kamino.py b/asv/benchmarks/simulation/bench_kamino.py index c180b0c77c..bad8d46ee5 100644 --- a/asv/benchmarks/simulation/bench_kamino.py +++ b/asv/benchmarks/simulation/bench_kamino.py @@ -14,7 +14,41 @@ parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) sys.path.append(parent_dir) -from benchmark_kamino import DRLegsBenchmarkWorkload +from benchmark_metrics import ( + _SimulationMetricTracks, + _SimulationMetricTracksUnparameterized, + collect_simulation_metrics, +) + + +def _collect_metrics_dr_legs(robot, world_count, num_frames, samples, use_policy): + if wp.get_cuda_device_count() == 0: + return None + + from benchmark_kamino import DRLegsBenchmarkWorkload # noqa: PLC0415 + + builder = DRLegsBenchmarkWorkload.create_model_builder(robot, world_count) + + def create_workload(): + workload = DRLegsBenchmarkWorkload( + robot=robot, + world_count=world_count, + use_cuda_graph=True, + use_policy=use_policy, + builder=builder, + ) + if workload.graph is None or workload.reset_graph is None: + raise RuntimeError("KPI benchmark requires CUDA graph capture (is the CUDA mempool allocator enabled?)") + wp.synchronize_device() + return workload + + return collect_simulation_metrics( + create_workload=create_workload, + world_count=world_count, + num_frames=num_frames, + samples=samples, + validate=lambda workload: workload.test_final(), + ) class _FastBenchmark: @@ -28,6 +62,8 @@ class _FastBenchmark: world_count = None def setup(self): + from benchmark_kamino import DRLegsBenchmarkWorkload # noqa: PLC0415 + if not hasattr(self, "_builder") or self._builder is None: self._builder = DRLegsBenchmarkWorkload.create_model_builder(self.robot, self.world_count) @@ -58,7 +94,7 @@ def time_simulate(self): wp.synchronize_device() -class _KpiBenchmark: +class _KpiBenchmark(_SimulationMetricTracks): """Utility base class for Kamino KPI benchmarks.""" param_names: ClassVar[list[str]] = ["world_count"] @@ -68,35 +104,20 @@ class _KpiBenchmark: samples = None use_policy = True - def setup(self, world_count): - if not hasattr(self, "_builder") or self._builder is None: - self._builder = {} - if world_count not in self._builder: - self._builder[world_count] = DRLegsBenchmarkWorkload.create_model_builder(self.robot, world_count) + def _collect_metrics(self): + if wp.get_cuda_device_count() == 0: + return None - @skip_benchmark_if(wp.get_cuda_device_count() == 0) - def track_simulate(self, world_count): - total_time = 0.0 - for _iter in range(self.samples): - workload = DRLegsBenchmarkWorkload( + metrics = {} + for world_count in self.params[0]: + metrics[world_count] = _collect_metrics_dr_legs( robot=self.robot, world_count=world_count, - use_cuda_graph=True, + num_frames=self.num_frames, + samples=self.samples, use_policy=self.use_policy, - builder=self._builder[world_count], ) - if workload.graph is None or workload.reset_graph is None: - raise RuntimeError("KPI benchmark requires CUDA graph capture (is the CUDA mempool allocator enabled?)") - - wp.synchronize_device() - for _ in range(self.num_frames): - workload.step() - total_time += workload.benchmark_time - workload.test_final() - - return total_time * 1000 / (self.num_frames * workload.sim_substeps * world_count * self.samples) - - track_simulate.unit = "ms/world-step" + return metrics class FastDRLegs(_FastBenchmark): @@ -106,12 +127,101 @@ class FastDRLegs(_FastBenchmark): world_count = 32 +class FastMetricsDRLegs(_SimulationMetricTracksUnparameterized): + num_frames = 25 + robot = "dr_legs" + samples = 2 + world_count = 32 + + def setup_cache(self): + return _collect_metrics_dr_legs( + robot=self.robot, + world_count=self.world_count, + num_frames=self.num_frames, + samples=self.samples, + use_policy=False, + ) + + class KpiDRLegs(_KpiBenchmark): params: ClassVar[list[list[int]]] = [[4096]] num_frames = 25 robot = "dr_legs" samples = 2 + def setup_cache(self): + return self._collect_metrics() + + setup_cache.timeout = 1200 + + +class NotifyDRLegs: + """Benchmark Kamino model notifications for 2048 DR Legs worlds.""" + + number = 10 + repeat = 7 + rounds = 1 + timeout = 3600 + world_count = 2048 + + def setup(self): + from benchmark_kamino import DRLegsBenchmarkWorkload # noqa: PLC0415 + + import newton # noqa: PLC0415 + + builder = DRLegsBenchmarkWorkload.create_model_builder("dr_legs", self.world_count) + model = builder.finalize(skip_validation_joints=True) + self._solver = newton.solvers.SolverKamino(model) + self._model_flags = newton.ModelFlags + for flag in ( + self._model_flags.MODEL_PROPERTIES, + self._model_flags.BODY_PROPERTIES, + self._model_flags.BODY_INERTIAL_PROPERTIES, + self._model_flags.SHAPE_PROPERTIES, + self._model_flags.JOINT_PROPERTIES, + self._model_flags.JOINT_DOF_PROPERTIES, + self._model_flags.ACTUATOR_PROPERTIES, + self._model_flags.ALL, + ): + self._solver.notify_model_changed(flag) + wp.synchronize_device() + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_notify_actuator_properties(self): + self._notify(self._model_flags.ACTUATOR_PROPERTIES) + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_notify_all(self): + self._notify(self._model_flags.ALL) + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_notify_body_inertial_properties(self): + self._notify(self._model_flags.BODY_INERTIAL_PROPERTIES) + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_notify_body_properties(self): + self._notify(self._model_flags.BODY_PROPERTIES) + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_notify_joint_dof_properties(self): + self._notify(self._model_flags.JOINT_DOF_PROPERTIES) + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_notify_joint_properties(self): + self._notify(self._model_flags.JOINT_PROPERTIES) + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_notify_model_properties(self): + self._notify(self._model_flags.MODEL_PROPERTIES) + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def time_notify_shape_properties(self): + self._notify(self._model_flags.SHAPE_PROPERTIES) + + def _notify(self, flag): + self._solver.notify_model_changed(flag) + wp.synchronize_device() + if __name__ == "__main__": import argparse @@ -120,7 +230,9 @@ class KpiDRLegs(_KpiBenchmark): benchmark_list = { "FastDRLegs": FastDRLegs, + "FastMetricsDRLegs": FastMetricsDRLegs, "KpiDRLegs": KpiDRLegs, + "NotifyDRLegs": NotifyDRLegs, } parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) diff --git a/asv/benchmarks/simulation/bench_mujoco.py b/asv/benchmarks/simulation/bench_mujoco.py index f449a73996..0c37eca685 100644 --- a/asv/benchmarks/simulation/bench_mujoco.py +++ b/asv/benchmarks/simulation/bench_mujoco.py @@ -1,25 +1,47 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers # SPDX-License-Identifier: Apache-2.0 +import math import os import sys +import time +from dataclasses import replace +from functools import partial +import numpy as np import warp as wp wp.config.enable_backward = False wp.config.log_level = wp.LOG_WARNING -from asv_runner.benchmarks.mark import skip_benchmark_if +from asv_runner.benchmarks.mark import SkipNotImplemented, skip_benchmark_if parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) sys.path.append(parent_dir) -from benchmark_mujoco import Example +from benchmark_metrics import ( + _SimulationMetricTracks, + collect_simulation_metrics, +) -from newton.utils import EventTracer +class _SimulationMetricTracksMuJoCo(_SimulationMetricTracks): + """MuJoCo-specific tracked metrics.""" -class _KpiBenchmark: + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_solver_niter_mean(self, metrics, world_count): + return metrics[world_count].solver_niter_mean + + track_solver_niter_mean.unit = "iterations" + + @skip_benchmark_if(wp.get_cuda_device_count() == 0) + def track_solver_niter_max(self, metrics, world_count): + return metrics[world_count].solver_niter_max + + track_solver_niter_max.unit = "iterations" + + +class _KpiBenchmark(_SimulationMetricTracksMuJoCo): """Utility base class for KPI benchmarks.""" param_names = ["world_count"] @@ -27,43 +49,163 @@ class _KpiBenchmark: params = None robot = None samples = None - ls_iteration = None random_init = None environment = "None" + expected_bodies_per_world = None + + def _create_workload(self, builder, world_count): + from benchmark_mujoco import Example # noqa: PLC0415 + + workload = Example( + stage_path=None, + robot=self.robot, + randomize=self.random_init, + headless=True, + actuation="random", + use_cuda_graph=True, + builder=builder, + world_count=world_count, + environment=self.environment, + ) + if workload.graph is None: + raise RuntimeError("KPI benchmark requires CUDA graph capture (is the CUDA mempool allocator enabled?)") + wp.synchronize_device() + return workload + + def _validate_workload(self, workload, world_count): + workload.test_final() + if self.expected_bodies_per_world is None: + return + expected_body_count = self.expected_bodies_per_world * world_count + if workload.model.body_count != expected_body_count: + raise RuntimeError( + f"Expected {self.expected_bodies_per_world} bodies per world for {self.environment}, " + f"got {workload.model.body_count / world_count:g}" + ) - def setup(self, world_count): - if not hasattr(self, "builder") or self.builder is None: - self.builder = {} - if world_count not in self.builder: - self.builder[world_count] = Example.create_model_builder( - self.robot, world_count, randomize=self.random_init, seed=123 + def _validate_metrics_workload(self, workload, world_count, solver_niter_samples): + self._validate_workload(workload, world_count) + solver_niter_samples.append(workload.solver.mjw_data.solver_niter.numpy()) + + def _collect_metrics(self): + if wp.get_cuda_device_count() == 0: + return None + + from benchmark_mujoco import Example # noqa: PLC0415 + + metrics = {} + for world_count in self.params[0]: + builder = Example.create_model_builder( + self.robot, + world_count, + environment=self.environment, + randomize=self.random_init, + seed=123, ) + solver_niter_samples = [] + + def create_workload(builder=builder, world_count=world_count): + return self._create_workload(builder, world_count) + + world_metrics = collect_simulation_metrics( + create_workload=create_workload, + world_count=world_count, + num_frames=self.num_frames, + samples=self.samples, + validate=partial( + self._validate_metrics_workload, + world_count=world_count, + solver_niter_samples=solver_niter_samples, + ), + ) + solver_niter = np.concatenate([np.asarray(values).reshape(-1) for values in solver_niter_samples]) + metrics[world_count] = replace( + world_metrics, + solver_niter_mean=float(np.mean(solver_niter)), + solver_niter_max=float(np.max(solver_niter)), + ) + return metrics - @skip_benchmark_if(wp.get_cuda_device_count() == 0) - def track_simulate(self, world_count): - total_time = 0.0 - for _iter in range(self.samples): - example = Example( + +class _RealtimePhysicsBenchmark: + """Report single-world physics throughput, stability, and real-time factor.""" + + robot = None + physics_hz = 200 + num_steps = 300 + warmup_steps = 60 + repeat = 3 + number = 1 + rounds = 2 + timeout = 600 + + def setup(self): + if wp.get_cuda_device_count() == 0: + raise SkipNotImplemented + + from benchmark_mujoco import Example # noqa: PLC0415 + + with wp.ScopedDevice("cuda:0"): + if not wp.is_mempool_enabled(wp.get_device()): + raise SkipNotImplemented + builder = Example.create_model_builder(self.robot, 1, randomize=True, seed=123) + self.example = Example( stage_path=None, robot=self.robot, - randomize=self.random_init, + randomize=True, headless=True, - actuation="random", + actuation="None", use_cuda_graph=True, - builder=self.builder[world_count], - ls_iteration=self.ls_iteration, - environment=self.environment, + builder=builder, + fps=self.physics_hz, + sim_substeps=1, ) + for _ in range(self.warmup_steps): + self.example.step() + + def track_mean_step_ms(self) -> float: + return 1000.0 * self._mean(self._measure_step_durations()) + + track_mean_step_ms.unit = "ms/step" + + def track_p95_step_ms(self) -> float: + durations = sorted(self._measure_step_durations()) + p95_index = min(len(durations) - 1, int(math.ceil(0.95 * len(durations))) - 1) + return 1000.0 * durations[p95_index] + + track_p95_step_ms.unit = "ms/step" + + def track_step_rate_hz(self) -> float: + return 1.0 / self._mean(self._measure_step_durations()) + + track_step_rate_hz.unit = "Hz" + + def track_step_time_cv_pct(self) -> float: + durations = self._measure_step_durations() + mean = self._mean(durations) + variance = sum((duration - mean) ** 2 for duration in durations) / len(durations) + return 100.0 * math.sqrt(variance) / mean + + track_step_time_cv_pct.unit = "%" + + def track_real_time_factor(self) -> float: + mean_step_s = self._mean(self._measure_step_durations()) + return self.example.sim_dt / mean_step_s - wp.synchronize_device() - for _ in range(self.num_frames): - example.step() - wp.synchronize_device() - total_time += example.benchmark_time + track_real_time_factor.unit = "x" - return total_time * 1000 / (self.num_frames * example.sim_substeps * world_count * self.samples) + def _measure_step_durations(self) -> list[float]: + durations = [] + with wp.ScopedDevice("cuda:0"): + for _ in range(self.num_steps): + start = time.perf_counter() + self.example.step() + durations.append(time.perf_counter() - start) + return durations - track_simulate.unit = "ms/world-step" + @staticmethod + def _mean(values: list[float]) -> float: + return sum(values) / len(values) class _NewtonOverheadBenchmark: @@ -74,10 +216,11 @@ class _NewtonOverheadBenchmark: params = None robot = None samples = None - ls_iteration = None random_init = None def setup(self, world_count): + from benchmark_mujoco import Example # noqa: PLC0415 + if not hasattr(self, "builder") or self.builder is None: self.builder = {} if world_count not in self.builder: @@ -87,6 +230,10 @@ def setup(self, world_count): @skip_benchmark_if(wp.get_cuda_device_count() == 0) def track_simulate(self, world_count): + from benchmark_mujoco import Example # noqa: PLC0415 + + from newton.utils import EventTracer # noqa: PLC0415 + trace = {} with EventTracer(enabled=True) as tracer: for _iter in range(self.samples): @@ -99,7 +246,6 @@ def track_simulate(self, world_count): world_count=world_count, use_cuda_graph=True, builder=self.builder[world_count], - ls_iteration=self.ls_iteration, ) for _ in range(self.num_frames): @@ -120,10 +266,12 @@ class FastCartpole(_KpiBenchmark): num_frames = 50 robot = "cartpole" samples = 4 - ls_iteration = 3 random_init = True environment = "None" + def setup_cache(self): + return self._collect_metrics() + class FastG1(_KpiBenchmark): params = [[8192]] @@ -131,10 +279,12 @@ class FastG1(_KpiBenchmark): robot = "g1" timeout = 900 samples = 2 - ls_iteration = 10 random_init = True environment = "None" + def setup_cache(self): + return self._collect_metrics() + class FastNewtonOverheadG1(_NewtonOverheadBenchmark): params = [[8192]] @@ -142,7 +292,6 @@ class FastNewtonOverheadG1(_NewtonOverheadBenchmark): robot = "g1" timeout = 900 samples = 2 - ls_iteration = 10 random_init = True @@ -151,17 +300,24 @@ class FastHumanoid(_KpiBenchmark): num_frames = 100 robot = "humanoid" samples = 4 - ls_iteration = 15 random_init = True environment = "None" + def setup_cache(self): + return self._collect_metrics() + + +class RealtimeHumanoidPhysics(_RealtimePhysicsBenchmark): + """Single highly articulated humanoid in physics-only mode.""" + + robot = "humanoid" + class FastNewtonOverheadHumanoid(_NewtonOverheadBenchmark): params = [[8192]] num_frames = 100 robot = "humanoid" samples = 4 - ls_iteration = 15 random_init = True @@ -171,20 +327,27 @@ class FastAllegro(_KpiBenchmark): robot = "allegro" timeout = 900 samples = 2 - ls_iteration = 10 random_init = False environment = "None" + def setup_cache(self): + return self._collect_metrics() + class FastKitchenG1(_KpiBenchmark): + # #3574 bounds replicated filter pairs to colliding shapes so 512 worlds fit on CI hosts. params = [[512]] num_frames = 50 robot = "g1" timeout = 900 + version = "2" # The pre-v2 series accidentally omitted the kitchen environment. samples = 2 - ls_iteration = 10 random_init = True environment = "kitchen" + expected_bodies_per_world = 111 + + def setup_cache(self): + return self._collect_metrics() if __name__ == "__main__": @@ -200,6 +363,7 @@ class FastKitchenG1(_KpiBenchmark): "FastKitchenG1": FastKitchenG1, "FastNewtonOverheadG1": FastNewtonOverheadG1, "FastNewtonOverheadHumanoid": FastNewtonOverheadHumanoid, + "RealtimeHumanoidPhysics": RealtimeHumanoidPhysics, } parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) diff --git a/asv/benchmarks/simulation/bench_quadruped_xpbd.py b/asv/benchmarks/simulation/bench_quadruped_xpbd.py index 2f1ce6e4b4..1f9ae001bf 100644 --- a/asv/benchmarks/simulation/bench_quadruped_xpbd.py +++ b/asv/benchmarks/simulation/bench_quadruped_xpbd.py @@ -1,15 +1,35 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers # SPDX-License-Identifier: Apache-2.0 +import os +import sys + import warp as wp from asv_runner.benchmarks.mark import skip_benchmark_if wp.config.enable_backward = False wp.config.log_level = wp.LOG_WARNING -import newton -import newton.examples -from newton.examples.basic.example_basic_urdf import Example +parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.append(parent_dir) + +from benchmark_metrics import ( + _SimulationMetricTracksUnparameterized, + collect_simulation_metrics, + validate_simulation_state, +) + + +def _create_example(num_frames, world_count): + import newton # noqa: PLC0415 + import newton.examples # noqa: PLC0415 + from newton.examples.basic.example_basic_urdf import Example # noqa: PLC0415 + + if hasattr(newton.examples, "default_args") and hasattr(Example, "create_parser"): + args = newton.examples.default_args(Example.create_parser()) + args.world_count = world_count + return Example(newton.viewer.ViewerNull(num_frames=num_frames), args) + return Example(newton.viewer.ViewerNull(num_frames=num_frames), world_count) class FastExampleQuadrupedXPBD: @@ -18,12 +38,7 @@ class FastExampleQuadrupedXPBD: def setup(self): self.num_frames = 1000 - if hasattr(newton.examples, "default_args") and hasattr(Example, "create_parser"): - args = newton.examples.default_args(Example.create_parser()) - args.world_count = 200 - self.example = Example(newton.viewer.ViewerNull(num_frames=self.num_frames), args) - else: - self.example = Example(newton.viewer.ViewerNull(num_frames=self.num_frames), 200) + self.example = _create_example(self.num_frames, world_count=200) @skip_benchmark_if(wp.get_cuda_device_count() == 0) def time_simulate(self): @@ -32,6 +47,33 @@ def time_simulate(self): wp.synchronize_device() +class FastMetricsExampleQuadrupedXPBD(_SimulationMetricTracksUnparameterized): + num_frames = 1000 + samples = 1 + world_count = 200 + + def setup_cache(self): + if wp.get_cuda_device_count() == 0: + return None + + def validate_workload(workload): + validate_simulation_state( + workload.state_0, + max_linear_speed=0.3, + max_angular_speed=0.3, + ) + workload.test_final() + + return collect_simulation_metrics( + create_workload=lambda: _create_example(self.num_frames, self.world_count), + world_count=self.world_count, + num_frames=self.num_frames, + samples=self.samples, + synchronize=wp.synchronize_device, + validate=validate_workload, + ) + + if __name__ == "__main__": import argparse @@ -39,6 +81,7 @@ def time_simulate(self): benchmark_list = { "FastExampleQuadrupedXPBD": FastExampleQuadrupedXPBD, + "FastMetricsExampleQuadrupedXPBD": FastMetricsExampleQuadrupedXPBD, } parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) diff --git a/asv/benchmarks/simulation/bench_teleop_mujoco.py b/asv/benchmarks/simulation/bench_teleop_mujoco.py new file mode 100644 index 0000000000..a1aa78cb5c --- /dev/null +++ b/asv/benchmarks/simulation/bench_teleop_mujoco.py @@ -0,0 +1,595 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark a scripted G1 bimanual pushing control loop. + +The measured loop covers deterministic two-hand six-DoF command generation, IK, +joint-target writes, and two completed MuJoCo physics substeps. Rendering, +physical input devices, transport, perception, and display latency are outside +the benchmark scope. +""" + +from __future__ import annotations + +import copy +import math +import time +from collections import defaultdict, deque +from dataclasses import dataclass +from typing import ClassVar + +import numpy as np +import warp as wp +from asv_runner.benchmarks.mark import SkipNotImplemented + +wp.config.enable_backward = False +wp.config.log_level = wp.LOG_WARNING + +import newton +import newton.ik as ik +import newton.utils +from newton import JointTargetMode + + +@wp.kernel +def _write_robot_targets( + joint_q_ik: wp.array2d[wp.float32], + nominal_joint_q: wp.array[wp.float32], + arm_joint_mask: wp.array[wp.int32], + previous_joint_target_q: wp.array[wp.float32], + dt: wp.float32, + joint_target_q: wp.array[wp.float32], + joint_target_qd: wp.array[wp.float32], +): + i = wp.tid() + q = nominal_joint_q[i] + if arm_joint_mask[i] != 0: + q = joint_q_ik[0, i] + + qd = 1.5 * (q - previous_joint_target_q[i]) / dt + joint_target_q[i] = q + joint_target_qd[i] = wp.clamp(qd, -20.0, 20.0) + previous_joint_target_q[i] = q + + +@dataclass(frozen=True) +class _TeleopMode: + device: str + mujoco_backend: str + use_graph: bool = False + requires_cuda: bool = False + requires_cuda_graph: bool = False + + +_TELEOP_MODES = { + "mjwarp_cuda_graph": _TeleopMode( + device="cuda:0", + mujoco_backend="warp", + use_graph=True, + requires_cuda=True, + requires_cuda_graph=True, + ), + "mjwarp_cpu_graph": _TeleopMode(device="cpu", mujoco_backend="warp", use_graph=True), + "mjwarp_cpu_eager": _TeleopMode(device="cpu", mujoco_backend="warp"), + "mujoco_cpu_eager": _TeleopMode(device="cpu", mujoco_backend="cpu"), +} + + +class _WindowStats: + def __init__(self, maxlen: int): + self.values: dict[str, deque[float]] = defaultdict(lambda: deque(maxlen=maxlen)) + + def add(self, name: str, value: float) -> None: + if math.isfinite(value): + self.values[name].append(float(value)) + + def summary(self, name: str) -> tuple[float, float, float]: + values = self.values.get(name) + if not values: + raise RuntimeError(f"No teleop samples collected for {name!r}") + ordered = sorted(values) + p95_index = min(len(ordered) - 1, math.ceil(0.95 * len(ordered)) - 1) + return sum(values) / len(values), ordered[p95_index], ordered[-1] + + def coefficient_of_variation_pct(self, name: str) -> float: + values = self.values.get(name) + if not values: + raise RuntimeError(f"No teleop samples collected for {name!r}") + mean = sum(values) / len(values) + if mean == 0.0: + return 0.0 + variance = sum((value - mean) ** 2 for value in values) / len(values) + return 100.0 * math.sqrt(variance) / mean + + def clear(self) -> None: + self.values.clear() + + +def _quat_to_vec4(q: wp.quat) -> wp.vec4: + return wp.vec4(float(q[0]), float(q[1]), float(q[2]), float(q[3])) + + +def _quat_to_np(q: wp.quat) -> np.ndarray: + return np.array([float(q[0]), float(q[1]), float(q[2]), float(q[3])], dtype=np.float32) + + +def _vec3_to_np(v: wp.vec3) -> np.ndarray: + return np.array([float(v[0]), float(v[1]), float(v[2])], dtype=np.float32) + + +class _TeleopLoop: + # Keep the control loop below the physics rate so every command advances + # multiple completed physics steps. Performance targets remain external. + control_hz = 100 + physics_hz = 200 + sim_substeps = 2 + linear_speed = 0.5 + sweep_half_period = 1.6 + arm_joint_indices = (*range(15, 22), *range(29, 36)) + + def __init__(self, mode: _TeleopMode, stats_window: int): + self.device = wp.get_device(mode.device) + self.frame_dt = 1.0 / self.control_hz + self.sim_dt = 1.0 / self.physics_hz + if not math.isclose(self.frame_dt, self.sim_substeps * self.sim_dt): + raise ValueError("The control period must contain an integer number of physics steps") + self.sim_time = 0.0 + self.frame_index = 0 + self.stats = _WindowStats(stats_window) + + robot = self._build_robot() + self.model_ik = copy.deepcopy(robot).finalize() + + scene = newton.ModelBuilder() + scene.add_builder(robot) + self.object_body_index, self.object_shape_index = self._add_scene(scene) + self.model = scene.finalize() + + self.state_0 = self.model.state() + self.state_1 = self.model.state() + self.control = self.model.control() + newton.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, self.state_0) + + state_ik = self.model_ik.state() + newton.eval_fk(self.model_ik, self.model_ik.joint_q, self.model_ik.joint_qd, state_ik) + self.ee_indices = [ + self._find_body(self.model_ik, "left_wrist_yaw_link"), + self._find_body(self.model_ik, "right_wrist_yaw_link"), + ] + body_q = state_ik.body_q.numpy() + self.target_tfs = [wp.transform(*body_q[index]) for index in self.ee_indices] + self.target_rotations = [wp.transform_get_rotation(target) for target in self.target_tfs] + self.initial_object_position = self.state_0.body_q.numpy()[self.object_body_index, :3].copy() + self.hand_shape_indices = { + shape_index + for shape_index, shape_label in enumerate(self.model.shape_label) + if "/left_hand_" in shape_label or "/right_hand_" in shape_label + } + + self.nominal_joint_q = wp.clone(self.model.joint_q) + self.arm_joint_mask = wp.array( + [int(i in self.arm_joint_indices) for i in range(self.model.joint_coord_count)], + dtype=wp.int32, + device=self.device, + ) + self.previous_joint_target_q = wp.empty(self.model.joint_coord_count, dtype=wp.float32, device=self.device) + self._setup_ik() + wp.copy(self.control.joint_target_q, self.model.joint_q) + wp.copy(self.previous_joint_target_q, self.control.joint_target_q) + self._write_targets() + + self.solver = newton.solvers.SolverMuJoCo( + self.model, + use_mujoco_cpu=mode.mujoco_backend == "cpu", + solver="newton", + integrator="implicitfast", + cone="pyramidal", + njmax=200, + nconmax=100, + ) + self.mjc_geom_to_newton_shape = self.solver.mjc_geom_to_newton_shape.numpy()[0] + self.contacts = None + if mode.mujoco_backend == "warp": + self.contacts = newton.Contacts(self.solver.get_max_contact_count(), 0) + + self.graph_ik = None + if mode.use_graph: + with wp.ScopedCapture(device=self.device) as capture: + self.ik_solver.step(self.joint_q_ik, self.joint_q_ik, iterations=16) + self.graph_ik = capture.graph + + self.graph_sim = None + if mode.use_graph and mode.mujoco_backend == "warp": + with wp.ScopedCapture(device=self.device) as capture: + self._simulate() + self.graph_sim = capture.graph + + def _build_robot(self) -> newton.ModelBuilder: + robot = newton.ModelBuilder() + newton.solvers.SolverMuJoCo.register_custom_attributes(robot) + robot.default_joint_cfg = newton.ModelBuilder.JointDofConfig(limit_ke=1.0e3, limit_kd=1.0e1, friction=1.0e-5) + robot.default_shape_cfg.ke = 1.0e3 + robot.default_shape_cfg.kd = 2.0e2 + robot.default_shape_cfg.kf = 1.0e3 + robot.default_shape_cfg.mu = 0.75 + robot.add_usd( + str(newton.utils.download_asset("unitree_g1") / "usd_structured" / "g1_29dof_with_hand_rev_1_0.usda"), + floating=False, + collapse_fixed_joints=True, + enable_self_collisions=False, + hide_collision_shapes=True, + skip_mesh_approximation=True, + ) + robot.approximate_meshes("bounding_box") + + for i in range(robot.joint_dof_count): + robot.joint_target_ke[i] = 2500.0 if i in self.arm_joint_indices else 500.0 + robot.joint_target_kd[i] = 120.0 if i in self.arm_joint_indices else 20.0 + robot.joint_target_mode[i] = int(JointTargetMode.POSITION_VELOCITY) + for i in self.arm_joint_indices: + robot.joint_effort_limit[i] = 200.0 + return robot + + @staticmethod + def _add_scene(builder: newton.ModelBuilder) -> tuple[int, int]: + builder.add_shape_box( + body=-1, + xform=wp.transform(wp.vec3(0.48, 0.0, 0.72), wp.quat_identity()), + hx=0.34, + hy=0.50, + hz=0.04, + cfg=newton.ModelBuilder.ShapeConfig(mu=0.8, kd=50.0), + ) + box_size = wp.vec3(0.10, 0.10, 0.10) + box_body = builder.add_body(xform=wp.transform(wp.vec3(0.48, 0.0, 0.81), wp.quat_identity())) + box_shape = builder.add_shape_box( + body=box_body, + hx=0.5 * box_size[0], + hy=0.5 * box_size[1], + hz=0.5 * box_size[2], + cfg=newton.ModelBuilder.ShapeConfig(mu=1.2, kd=80.0, density=12000.0), + ) + builder.add_ground_plane() + return box_body, box_shape + + @staticmethod + def _find_body(model: newton.Model, name: str) -> int: + for index, label in enumerate(model.body_label): + if label.endswith(f"/{name}") or label == name: + return index + raise RuntimeError(f"Body {name!r} was not found in the teleop model") + + def _setup_ik(self) -> None: + self.pos_objectives = [] + self.rot_objectives = [] + for link_index, target in zip(self.ee_indices, self.target_tfs, strict=True): + target_pos = wp.transform_get_translation(target) + target_rot = wp.transform_get_rotation(target) + self.pos_objectives.append( + ik.IKObjectivePosition( + link_index=link_index, + link_offset=wp.vec3(0.0, 0.0, 0.0), + target_positions=wp.array([target_pos], dtype=wp.vec3, device=self.device), + ) + ) + self.rot_objectives.append( + ik.IKObjectiveRotation( + link_index=link_index, + link_offset_rotation=wp.quat_identity(), + target_rotations=wp.array([_quat_to_vec4(target_rot)], dtype=wp.vec4, device=self.device), + ) + ) + + nominal_q = self.model_ik.joint_q.numpy() + limit_lower = self.model_ik.joint_limit_lower.numpy() + limit_upper = self.model_ik.joint_limit_upper.numpy() + fixed_joint_indices = set(range(self.model_ik.joint_coord_count)) - set(self.arm_joint_indices) + for index in fixed_joint_indices: + limit_lower[index] = nominal_q[index] - 1.0e-4 + limit_upper[index] = nominal_q[index] + 1.0e-4 + joint_limit_objective = ik.IKObjectiveJointLimit( + joint_limit_lower=wp.array(limit_lower, device=self.device), + joint_limit_upper=wp.array(limit_upper, device=self.device), + weight=100.0, + ) + self.joint_q_ik = wp.array(self.model_ik.joint_q, shape=(1, self.model_ik.joint_coord_count)) + self.ik_solver = ik.IKSolver( + model=self.model_ik, + n_problems=1, + objectives=[*self.pos_objectives, *self.rot_objectives, joint_limit_objective], + lambda_initial=0.1, + jacobian_mode=ik.IKJacobianType.ANALYTIC, + ) + + def _update_command(self) -> None: + cycle_position = self.sim_time / self.sweep_half_period + active_hand = int(cycle_position) % 2 + progress = cycle_position % 1.0 + desired_positions = [ + np.array([0.24, 0.20, 0.95], dtype=np.float32), + np.array([0.24, -0.20, 0.95], dtype=np.float32), + ] + + sweep_start = 0.20 if active_hand == 0 else -0.20 + sweep_end = -0.20 if active_hand == 0 else 0.20 + active_position = desired_positions[active_hand] + # Advance above the box so the approach does not add a forward impulse. + if progress < 0.2: + active_position[0] = 0.24 + (0.38 - 0.24) * progress / 0.2 + elif progress < 0.3: + active_position[0] = 0.38 + active_position[2] = 0.95 + (0.84 - 0.95) * (progress - 0.2) / 0.1 + elif progress < 0.8: + sweep_progress = (progress - 0.3) / 0.5 + active_position[0] = 0.38 + active_position[1] = sweep_start + (sweep_end - sweep_start) * sweep_progress + active_position[2] = 0.84 + elif progress < 0.9: + active_position[0] = 0.38 + active_position[1] = sweep_end + active_position[2] = 0.84 + (0.95 - 0.84) * (progress - 0.8) / 0.1 + else: + active_position[0] = 0.38 + (0.24 - 0.38) * (progress - 0.9) / 0.1 + active_position[1] = sweep_end + + max_delta = self.linear_speed * self.frame_dt + for index, (target, desired_position) in enumerate(zip(self.target_tfs, desired_positions, strict=True)): + current_position = _vec3_to_np(wp.transform_get_translation(target)) + current_position += np.clip(desired_position - current_position, -max_delta, max_delta) + self.target_tfs[index] = wp.transform(wp.vec3(*current_position), self.target_rotations[index]) + + def _write_targets(self) -> None: + wp.launch( + _write_robot_targets, + dim=self.model_ik.joint_coord_count, + inputs=[ + self.joint_q_ik, + self.nominal_joint_q, + self.arm_joint_mask, + self.previous_joint_target_q, + self.frame_dt, + ], + outputs=[self.control.joint_target_q, self.control.joint_target_qd], + device=self.device, + ) + + def _simulate(self) -> None: + for _ in range(self.sim_substeps): + self.state_0.clear_forces() + self.state_1.clear_forces() + self.solver.step(self.state_0, self.state_1, self.control, None, self.sim_dt) + self.state_0, self.state_1 = self.state_1, self.state_0 + + def _record_workload_state(self) -> None: + body_q = self.state_0.body_q.numpy() + body_qd = self.state_0.body_qd.numpy() + if not np.isfinite(body_q).all() or not np.isfinite(body_qd).all(): + raise RuntimeError("The G1 pushing workload produced non-finite body state") + + position_errors = [] + rotation_errors = [] + for ee_index, target in zip(self.ee_indices, self.target_tfs, strict=True): + end_effector = body_q[ee_index] + target_pos = _vec3_to_np(wp.transform_get_translation(target)) + position_errors.append(float(np.linalg.norm(target_pos - end_effector[:3]))) + target_rot = _quat_to_np(wp.transform_get_rotation(target)) + quat_dot = min(1.0, abs(float(np.dot(end_effector[3:7], target_rot)))) + rotation_errors.append(2.0 * math.acos(quat_dot)) + self.stats.add("target_error_m", float(np.mean(position_errors))) + self.stats.add("target_rotation_error_rad", float(np.mean(rotation_errors))) + + object_position = body_q[self.object_body_index, :3] + self.stats.add("object_displacement_m", float(np.linalg.norm(object_position - self.initial_object_position))) + self.stats.add("hand_object_contact", float(self._has_hand_object_contact())) + + def _has_hand_object_contact(self) -> bool: + if self.contacts is not None: + self.solver.update_contacts(self.contacts, self.state_0) + contact_count = int(self.contacts.rigid_contact_count.numpy()[0]) + shape0 = self.contacts.rigid_contact_shape0.numpy()[:contact_count] + shape1 = self.contacts.rigid_contact_shape1.numpy()[:contact_count] + else: + contact_count = int(self.solver.mj_data.ncon) + geom_pairs = self.solver.mj_data.contact.geom[:contact_count] + shape0 = self.mjc_geom_to_newton_shape[geom_pairs[:, 0]] + shape1 = self.mjc_geom_to_newton_shape[geom_pairs[:, 1]] + + for first_shape, second_shape in zip(shape0, shape1, strict=True): + if (first_shape == self.object_shape_index and second_shape in self.hand_shape_indices) or ( + second_shape == self.object_shape_index and first_shape in self.hand_shape_indices + ): + return True + return False + + def step(self, *, collect_workload_state: bool = False) -> None: + frame_start = time.perf_counter() + self._update_command() + for pos_objective, rot_objective, target in zip( + self.pos_objectives, self.rot_objectives, self.target_tfs, strict=True + ): + pos_objective.set_target_position(0, wp.transform_get_translation(target)) + rot_objective.set_target_rotation(0, _quat_to_vec4(wp.transform_get_rotation(target))) + if self.graph_ik is None: + self.ik_solver.step(self.joint_q_ik, self.joint_q_ik, iterations=16) + else: + wp.capture_launch(self.graph_ik) + self._write_targets() + if self.graph_sim is None: + self._simulate() + else: + wp.capture_launch(self.graph_sim) + + if collect_workload_state: + self._record_workload_state() + else: + wp.synchronize_device(self.device) + self.stats.add("local_loop_ms", (time.perf_counter() - frame_start) * 1000.0) + + self.sim_time += self.frame_dt + self.frame_index += 1 + + def clear_metrics(self) -> None: + self.stats.clear() + + +def _skip_unavailable_mode(mode: _TeleopMode) -> None: + if mode.requires_cuda and wp.get_cuda_device_count() == 0: + raise SkipNotImplemented + if mode.requires_cuda_graph: + with wp.ScopedDevice(mode.device): + if not wp.is_mempool_enabled(wp.get_device()): + raise SkipNotImplemented + + +class _TeleopMuJoCoBenchmark: + """Shared setup for scripted synchronous teleop benchmarks.""" + + params: ClassVar[tuple[tuple[str, ...]]] = (tuple(_TELEOP_MODES.keys()),) + param_names: ClassVar[list[str]] = ["mode"] + repeat = 3 + number = 1 + rounds = 2 + timeout = 600 + num_frames = 300 + warmup_frames = 60 + + def setup(self, mode: str) -> None: + self.mode = _TELEOP_MODES[mode] + _skip_unavailable_mode(self.mode) + previous_target_layout = newton.use_coord_layout_targets + newton.use_coord_layout_targets = True + try: + with wp.ScopedDevice(self.mode.device): + self.loop = _TeleopLoop(self.mode, self.num_frames) + finally: + newton.use_coord_layout_targets = previous_target_layout + + self._step_frames(self.warmup_frames) + self.loop.clear_metrics() + + def time_teleop_loop(self, mode: str) -> None: + self._step_frames(self.num_frames) + + def _step_frames(self, frame_count: int, *, collect_workload_state: bool = False) -> None: + with wp.ScopedDevice(self.mode.device): + for _ in range(frame_count): + self.loop.step(collect_workload_state=collect_workload_state) + + def _measure_frames(self, *, collect_workload_state: bool = False) -> None: + self.loop.clear_metrics() + self._step_frames(self.num_frames, collect_workload_state=collect_workload_state) + + def _validate_workload(self) -> None: + if self.loop.stats.summary("hand_object_contact")[2] == 0.0: + raise RuntimeError("The G1 pushing workload did not produce hand-object contact") + if self.loop.stats.summary("object_displacement_m")[2] < 0.01: + raise RuntimeError("The G1 pushing workload did not move the object by at least 0.01 m") + + +class FastTeleopMuJoCo(_TeleopMuJoCoBenchmark): + """Pull-request smoke benchmarks across GPU and CPU execution modes.""" + + params: ClassVar[tuple[tuple[str, ...]]] = (tuple(_TELEOP_MODES),) + repeat = 2 + num_frames = 120 + warmup_frames = 30 + + def track_mean_loop_ms(self, mode: str) -> float: + self._measure_frames() + return self.loop.stats.summary("local_loop_ms")[0] + + track_mean_loop_ms.unit = "ms/frame" + + def track_p95_loop_ms(self, mode: str) -> float: + self._measure_frames() + return self.loop.stats.summary("local_loop_ms")[1] + + track_p95_loop_ms.unit = "ms/frame" + + +class TeleopMuJoCo(_TeleopMuJoCoBenchmark): + """Nightly teleop benchmark covering GPU and CPU solver backends.""" + + def track_mean_loop_ms(self, mode: str) -> float: + self._measure_frames() + return self.loop.stats.summary("local_loop_ms")[0] + + track_mean_loop_ms.unit = "ms/frame" + + def track_p95_loop_ms(self, mode: str) -> float: + self._measure_frames() + return self.loop.stats.summary("local_loop_ms")[1] + + track_p95_loop_ms.unit = "ms/frame" + + def track_frame_overrun_pct(self, mode: str) -> float: + self._measure_frames() + values = self.loop.stats.values["local_loop_ms"] + overruns = sum(value > self.loop.frame_dt * 1000.0 for value in values) + return 100.0 * overruns / len(values) + + track_frame_overrun_pct.unit = "%" + + def track_loop_time_cv_pct(self, mode: str) -> float: + self._measure_frames() + return self.loop.stats.coefficient_of_variation_pct("local_loop_ms") + + track_loop_time_cv_pct.unit = "%" + + def track_real_time_factor(self, mode: str) -> float: + self._measure_frames() + mean_loop_ms = self.loop.stats.summary("local_loop_ms")[0] + return self.loop.frame_dt * 1000.0 / mean_loop_ms + + track_real_time_factor.unit = "x" + + def track_sustainable_physics_step_hz(self, mode: str) -> float: + self._measure_frames() + mean_loop_ms = self.loop.stats.summary("local_loop_ms")[0] + return self.loop.sim_substeps * 1000.0 / mean_loop_ms + + track_sustainable_physics_step_hz.unit = "Hz" + + def track_mean_target_error_m(self, mode: str) -> float: + self._measure_frames(collect_workload_state=True) + return self.loop.stats.summary("target_error_m")[0] + + track_mean_target_error_m.unit = "m" + + def track_mean_target_rotation_error_rad(self, mode: str) -> float: + self._measure_frames(collect_workload_state=True) + return self.loop.stats.summary("target_rotation_error_rad")[0] + + track_mean_target_rotation_error_rad.unit = "rad" + + def track_hand_object_contact_frame_pct(self, mode: str) -> float: + self._measure_frames(collect_workload_state=True) + self._validate_workload() + return 100.0 * self.loop.stats.summary("hand_object_contact")[0] + + track_hand_object_contact_frame_pct.unit = "%" + + def track_object_displacement_m(self, mode: str) -> float: + self._measure_frames(collect_workload_state=True) + self._validate_workload() + return self.loop.stats.summary("object_displacement_m")[2] + + track_object_displacement_m.unit = "m" + + +if __name__ == "__main__": + import argparse + + from newton.utils import run_benchmark + + benchmark_list = { + "FastTeleopMuJoCo": FastTeleopMuJoCo, + "TeleopMuJoCo": TeleopMuJoCo, + } + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("-b", "--bench", action="append", choices=benchmark_list.keys()) + args = parser.parse_known_args()[0] + + benchmarks = args.bench if args.bench is not None else benchmark_list.keys() + for key in benchmarks: + run_benchmark(benchmark_list[key]) diff --git a/docs/api/newton_geometry.rst b/docs/api/newton_geometry.rst index 60043b88d4..e5900e28ed 100644 --- a/docs/api/newton_geometry.rst +++ b/docs/api/newton_geometry.rst @@ -25,8 +25,6 @@ newton.geometry :toctree: _generated :signatures: long - build_bvh_particle - build_bvh_shape collide_box_box collide_capsule_box collide_capsule_capsule @@ -42,8 +40,6 @@ newton.geometry compute_inertia_shape compute_offset_mesh create_empty_sdf_data - refit_bvh_particle - refit_bvh_shape sdf_box sdf_capsule sdf_cone diff --git a/docs/api/newton_solvers.rst b/docs/api/newton_solvers.rst index b9ad1e59a3..20f79dc84d 100644 --- a/docs/api/newton_solvers.rst +++ b/docs/api/newton_solvers.rst @@ -39,7 +39,6 @@ https://newton-physics.github.io/newton/stable/solvers/index.html. SolverImplicitMPM SolverKamino SolverMuJoCo - SolverNotifyFlags SolverSemiImplicit SolverStyle3D SolverVBD diff --git a/docs/api/newton_utils.rst b/docs/api/newton_utils.rst index e5ade25c39..18f41d240a 100644 --- a/docs/api/newton_utils.rst +++ b/docs/api/newton_utils.rst @@ -13,6 +13,7 @@ newton.utils :toctree: _generated :nosignatures: + CableStiffness ColorSpace EventTracer MeshAdjacency @@ -38,6 +39,7 @@ newton.utils load_texture normalize_texture plot_graph + rasterize_mesh_to_heightfield remesh_mesh run_benchmark solidify_mesh diff --git a/docs/concepts/sensors.rst b/docs/concepts/sensors.rst index 7b0e4997bb..6797b03728 100644 --- a/docs/concepts/sensors.rst +++ b/docs/concepts/sensors.rst @@ -79,13 +79,25 @@ support label matching accept one of the following: - A **list of integer indices** -- selects directly by index. - A **single string pattern** -- selects all entries whose label matches the pattern via :func:`fnmatch.fnmatch` (supports ``*`` and ``?`` wildcards). -- A **list of string patterns** -- selects all entries whose label matches at least one of the patterns. +- A **list of string patterns** -- selects all entries whose label matches at least one pattern. +- A **compiled string regular expression** -- selects all entries whose entire label or name matches the expression via + :meth:`re.Pattern.fullmatch`. -Examples:: +Ordinary strings always use glob syntax. Compile a pattern with :func:`re.compile` to opt into regular-expression +syntax. Callers who want a regular expression to match a substring can add ``.*`` around that substring explicitly. +For :class:`~newton.selection.ArticulationView`, ``pattern`` is matched against full articulation labels. Joint and +link filters are matched against the final path component of each label. + +.. code-block:: python + + import re # single pattern: all shapes whose label starts with "foot_" SensorIMU(model, sites="foot_*") + # compiled regular expression: full-match an environment and object label + SensorIMU(model, sites=re.compile(r"/World/envs/env_[0-9]+/imu_(left|right)")) + # list of patterns: union of two groups SensorContact(model, sensing_shapes=["*Plate*", "*Flap*"]) diff --git a/docs/concepts/simulation_tuning_solvers.rst b/docs/concepts/simulation_tuning_solvers.rst index 85731e0e1a..c12b4d7b10 100644 --- a/docs/concepts/simulation_tuning_solvers.rst +++ b/docs/concepts/simulation_tuning_solvers.rst @@ -138,9 +138,7 @@ repository examples spend tuning effort, not a shared solver API. ``rigid_avbd_beta``, ``rigid_avbd_linear_beta``, ``rigid_avbd_angular_beta``, ``rigid_avbd_gamma``, ``rigid_contact_hard``, ``rigid_contact_history``, - ``rigid_contact_k_start``, ``rigid_contact_stick_motion_eps``, - ``rigid_contact_stick_freeze_translation_eps``, - ``rigid_contact_stick_freeze_angular_eps``, + ``rigid_contact_k_start``, ``rigid_body_contact_buffer_size``, ``rigid_body_particle_contact_buffer_size``, ``rigid_joint_linear_ke``, ``rigid_joint_angular_ke``, @@ -157,15 +155,22 @@ repository examples spend tuning effort, not a shared solver API. ``particle_topological_contact_filter_threshold``, ``particle_rest_shape_contact_exclusion_radius``. - Contact history requires matched contacts, for example - ``CollisionPipeline(contact_matching="latest")``. When recording VBD - steps in a CUDA graph, construct :class:`~newton.CollisionPipeline` - before :class:`~newton.solvers.SolverVBD` so contact history is - pre-allocated, or run one uncaptured solver step before capture. Buffer - sizes that are too small can drop contacts; sizes that are too large cost - memory and performance. Examples commonly tune ``iterations``, particle - self-contact radius and margin, particle contact buffers and filters, - ``particle_collision_detection_interval``, ``particle_enable_tile_solve``, - ``rigid_body_contact_buffer_size``, + ``CollisionPipeline(contact_matching="latest")`` or ``"sticky"``. + SolverVBD uses match indices only for numeric warm-starting; contact + geometry remains owned by the collision pipeline. Contact history is + cross-replay-persistent state, so it must always be pre-allocated + before graph capture on any device; otherwise SolverVBD raises a + ``RuntimeError``. Construct :class:`~newton.CollisionPipeline` before + :class:`~newton.solvers.SolverVBD` so contact history is pre-allocated, + or run one uncaptured solver step before capture. Ordinary contact + buffers can still grow on demand during graph capture on CPU and on + CUDA with the memory pool enabled; only CUDA capture without a memory + pool requires that they also be pre-allocated. Buffer sizes that are + too small can drop contacts; sizes that are too large cost memory and + performance. Examples commonly tune + ``iterations``, particle self-contact radius and margin, particle + contact buffers and filters, ``particle_collision_detection_interval``, + ``particle_enable_tile_solve``, ``rigid_body_contact_buffer_size``, ``rigid_body_particle_contact_buffer_size``, ``rigid_contact_hard``, ``rigid_contact_history``, and ``rigid_avbd_contact_alpha``. * - :class:`~newton.solvers.SolverFeatherstone` diff --git a/docs/concepts/worlds.rst b/docs/concepts/worlds.rst index 2c0d843a7c..50939f44de 100644 --- a/docs/concepts/worlds.rst +++ b/docs/concepts/worlds.rst @@ -266,6 +266,33 @@ While :meth:`~newton.ModelBuilder.begin_world` and :meth:`~newton.ModelBuilder.e world_count: 4 +.. _implicit-mpm-worlds: + +Implicit MPM world isolation +---------------------------- + +.. experimental:: + + Independent per-world Implicit MPM simulation and collider filtering may + change without prior notice. + +A multi-world :class:`~newton.solvers.SolverImplicitMPM` uses one shared FEM +topology by default, so world assignment alone does not isolate overlapping MPM +particles. Set :attr:`~newton.solvers.SolverImplicitMPM.Config.separate_worlds` +to create one FEM environment per world and isolate grid mass, momentum, +stress, and collider response. + +Isolated MPM requires every particle to belong to a local world. World-local +colliders affect only that world, while global static colliders and global +colliders backed by kinematic bodies affect every world. In isolated mode, +global colliders backed by dynamic bodies are rejected. + +See :class:`~newton.solvers.SolverImplicitMPM` for sparse-grid capacity and +CUDA graph-capture requirements. See +:meth:`~newton.solvers.SolverImplicitMPM.setup_collider` and +:meth:`~newton.solvers.SolverImplicitMPM.reset` for collider configuration and +selective reset behavior. + .. _Per-world gravity: Per-World Gravity diff --git a/docs/guide/development.rst b/docs/guide/development.rst index 79433dc5ac..0215c0770f 100644 --- a/docs/guide/development.rst +++ b/docs/guide/development.rst @@ -799,6 +799,43 @@ benchmark code from the ``asv/benchmarks`` directory against the code state of t the benchmark definitions themselves are not checked out from different branches—only the code being benchmarked is. +Simulation benchmark metrics +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The nightly robot and policy simulation benchmarks publish several metrics from one cached benchmark run. This +includes the batched MuJoCo and Kamino KPI workloads, the non-policy DR Legs workload, the pretrained Anymal +workload, and the XPBD quadruped workload. Companion ``FastMetrics*`` classes reuse the exact factories, scenes, +solvers, and frame counts of existing aggregate benchmarks. Rendering, IK, inverse dynamics, CPU-backend +regression, material/contact microbenchmarks, and startup benchmarks are not experience-collection workloads and +are outside this metric set. + +Let ``F`` be the number of measured frames across all samples, ``S`` the number of physics substeps per frame, +``W`` the world count, ``dt`` the physics timestep in seconds, ``T`` the synchronized complete-step wall time in +seconds, and ``T_physics`` the workload's synchronized internal physics time. The reported metrics are: + +* mean world-step time: ``1000 * T_physics / (F * S * W)`` in ``ms/world-step`` for workloads with an internal + timer and ``1000 * T / (F * S * W)`` otherwise; +* simulation throughput: ``F * S * W / T`` in ``world-steps/s``; +* real-time factor: ``F * S * W * dt / T``; +* p95 step time: the 95th percentile of synchronized complete-step times in ``ms/frame``; and +* steady-state GPU memory in ``MiB``; and +* mean and maximum MuJoCo solver iteration counts across worlds at the final frame of each measured sample. + +The existing KPI mean world-step series keeps its ``track_simulate`` name and definition, and existing +``time_simulate`` aggregate series remain unchanged. ``world_count`` remains an ASV parameter where applicable, +while ``sim_dt`` and ``sim_substeps`` are also recorded as tracked values so dashboard results retain the +configuration needed to interpret throughput and real-time factor. The existing KPI mean series continues to use +each workload's internal physics timer. Throughput, real-time factor, and p95 metrics time and synchronize the +complete ``step()`` operation, including any policy or control work it performs. + +GPU memory is the decrease in ``Device.free_memory`` between a baseline immediately before the first finalized +workload is created and a measurement after initialization, CUDA graph capture, and the first complete measured +sample, while that sample's workload is still live. This device-level delta includes allocations from Warp, +PyTorch, solver support, and CUDA graphs. Because the measurement is device-wide, runners must provide exclusive +GPU access during the measurement interval. The remaining samples do not affect this measurement. A benchmark +fails instead of publishing metrics if its final simulation state is invalid, has non-normalized body rotations, +or exceeds the workload's body-speed bounds. + Benchmarks can also be run against a range of commits using the ``commit1..commit2`` syntax. This is useful for comparing performance across several recent changes: diff --git a/docs/images/examples/example_balance_bird.jpg b/docs/images/examples/example_balance_bird.jpg new file mode 100644 index 0000000000..5ae44c4781 Binary files /dev/null and b/docs/images/examples/example_balance_bird.jpg differ diff --git a/docs/images/examples/example_cable_plectoneme.jpg b/docs/images/examples/example_cable_plectoneme.jpg new file mode 100644 index 0000000000..c20c61935f Binary files /dev/null and b/docs/images/examples/example_cable_plectoneme.jpg differ diff --git a/docs/images/examples/example_domino_spiral.jpg b/docs/images/examples/example_domino_spiral.jpg new file mode 100644 index 0000000000..1c1c418f4c Binary files /dev/null and b/docs/images/examples/example_domino_spiral.jpg differ diff --git a/docs/images/examples/example_newton_cradle.jpg b/docs/images/examples/example_newton_cradle.jpg new file mode 100644 index 0000000000..8cbf0bd744 Binary files /dev/null and b/docs/images/examples/example_newton_cradle.jpg differ diff --git a/docs/solvers/index.rst b/docs/solvers/index.rst index c08fed73bf..1ff9ebd489 100644 --- a/docs/solvers/index.rst +++ b/docs/solvers/index.rst @@ -176,7 +176,9 @@ formulation. :class:`~newton.solvers.SolverMuJoCo`, and :class:`~newton.solvers.SolverVBD`. - ``kf`` / ``ka``: :class:`~newton.solvers.SolverFeatherstone` and - :class:`~newton.solvers.SolverSemiImplicit`. + :class:`~newton.solvers.SolverSemiImplicit`; ``kf`` is also used by + :class:`~newton.solvers.SolverMuJoCo` + (see :ref:`mujoco-contact-friction-solreffriction`). - ``restitution``: :class:`~newton.solvers.SolverXPBD` when ``enable_restitution=True``, and :class:`~newton.solvers.SolverKamino`. - ``mu_torsional`` / ``mu_rolling``: :class:`~newton.solvers.SolverXPBD` and diff --git a/docs/solvers/kamino.rst b/docs/solvers/kamino.rst index 19c0124ee1..2785152532 100644 --- a/docs/solvers/kamino.rst +++ b/docs/solvers/kamino.rst @@ -24,3 +24,50 @@ primary requirements and an experimental solver is acceptable. See the :class:`~newton.solvers.SolverKamino` API reference for construction and configuration details. Runnable workflows are available in the `Kamino examples `_. + +Choosing a dynamics solver +-------------------------- + +Kamino provides two forward-dynamics backends: + +* ``"padmm"`` (default): proximal ADMM, dense Jacobians/dynamics, and the Euler + integrator. It is the slower, more robust option because it solves equality + and inequality constraints together. +* ``"dvi"`` (opt-in): projected dual iterations, sparse Jacobians, dense dynamics + with the RCM-reordered blocked LLT solver, and the Euler integrator. It is + generally faster, but approximates the coupled problem by alternating between + a direct solve for equality constraints and projected iterations for + inequality constraints. As a rule of thumb, DVI solves inequality constraints + less accurately than PADMM, particularly as the number of active inequalities + grows. Dual preconditioning is not supported. + +Select the backend when constructing the configuration so dependent defaults +initialize consistently: + +.. code-block:: python + + config = newton.solvers.SolverKamino.Config(dynamics_solver="dvi") + solver = newton.solvers.SolverKamino(model, config=config) + +DVI is best suited to performance-sensitive rigid mechanisms with relatively +few active contacts; PADMM remains the safer and more broadly validated choice. +Set ``sparse_jacobian=False`` for fully dense DVI, or set +``sparse_dynamics=True`` to use sparse dynamics with the Conjugate Residual +solver. With +``collect_solver_info=True``, DVI stores terminal residual status that should +not be interpreted as PADMM ADMM residuals. + +For large bilateral systems, opt into RCM-reordered factorization explicitly: + +.. code-block:: python + + config.dvi.bilateral_solver_type = "LLTBRCM" + config.dvi.bilateral_solver_kwargs = { + "block_size": 32, + "reuse_permutation": True, + "parallel_factorization": True, + } + +The cached permutation remains mathematically valid when matrix values or +sparsity change and is recomputed automatically if the active dimension +changes. Keep the default ``"LLTB"`` solver for small systems. diff --git a/docs/solvers/mujoco.rst b/docs/solvers/mujoco.rst index cb22b0f862..d674ba98d0 100644 --- a/docs/solvers/mujoco.rst +++ b/docs/solvers/mujoco.rst @@ -281,6 +281,35 @@ the mode from user code. For parameter interpretation, stability tradeoffs, and task-oriented guidance, see :ref:`Tuning MuJoCo`. +.. _mujoco-contact-friction-solreffriction: + +Contact friction ``solreffriction`` mapping +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For :class:`~newton.solvers.SolverMuJoCo`, ``kf`` maps to MuJoCo's per-contact +``solreffriction`` when the MuJoCo Warp backend uses elliptic friction cones +(``use_mujoco_cpu=False``, ``cone="elliptic"``) with Newton contacts +(``use_mujoco_contacts=False``). It targets the force-space friction slope +``f = -kf * v`` below the Coulomb limit. The mapping is exact when the sum of +MuJoCo's translational ``body_invweight0`` values matches the contact's inverse +effective mass (the relevant diagonal of :math:`J M^{-1} J^T`) and the contact +operates at its maximum impedance ``dmax``. Very large ``kf`` saturates at +MuJoCo's refsafe stability bound, where the reference time constant is clamped +to twice the timestep. + +The two shapes' ``kf`` values combine with the usual priority/``solmix`` +weighting. A resolved ``kf = 0`` makes the contact frictionless +(``condim = 1``), removing its sliding, torsional, and rolling friction rows. +If a positive ``kf`` cannot produce a positive, finite inverse-weight +denominator, ``solreffriction`` remains unset and MuJoCo inherits the normal +``solref``. The mapping is independent of the shape's ``solref_mode`` above, +which only governs the normal-direction ``solref``. + +The slope is calibrated for the sliding friction rows. With ``condim > 3``, the +torsional and rolling rows share the same per-contact ``solreffriction`` and +MuJoCo scales their regularization by the corresponding friction-coefficient +ratios, so their effective damping deviates from ``kf`` accordingly. + Actuators --------- diff --git a/newton/_src/geometry/__init__.py b/newton/_src/geometry/__init__.py index 00aa5c5026..e36ffa12f2 100644 --- a/newton/_src/geometry/__init__.py +++ b/newton/_src/geometry/__init__.py @@ -4,12 +4,6 @@ from .broad_phase_common import test_group_pair, test_world_and_group_pair from .broad_phase_nxn import BroadPhaseAllPairs, BroadPhaseExplicit from .broad_phase_sap import BroadPhaseSAP -from .bvh import ( - build_bvh_particle, - build_bvh_shape, - refit_bvh_particle, - refit_bvh_shape, -) from .collision_primitive import ( collide_box_box, collide_capsule_box, @@ -50,8 +44,6 @@ "ParticleFlags", "ShapeFlags", "TetMesh", - "build_bvh_particle", - "build_bvh_shape", "collide_box_box", "collide_capsule_box", "collide_capsule_capsule", @@ -69,8 +61,6 @@ "compute_shape_radius", "create_mesh_heightfield", "create_mesh_terrain", - "refit_bvh_particle", - "refit_bvh_shape", "test_group_pair", "test_world_and_group_pair", "transform_inertia", diff --git a/newton/_src/geometry/_deprecated_sdf_block_coords.py b/newton/_src/geometry/_deprecated_sdf_block_coords.py deleted file mode 100644 index f8e6a0d2b9..0000000000 --- a/newton/_src/geometry/_deprecated_sdf_block_coords.py +++ /dev/null @@ -1,99 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers -# SPDX-License-Identifier: Apache-2.0 - -"""Deprecated helpers for legacy SDF block-coordinate arrays. - -.. deprecated:: 1.3 - This entire module is deprecated and will be removed in a future release. - -The hydroelastic broadphase used to consume two precomputed arrays on -``Model`` — ``sdf_block_coords`` (flat ``wp.vec3us`` per active block) -and ``sdf_index2blocks`` (``[start, end)`` per SDF). Both were dropped -when the broadphase began deriving block coordinates arithmetically -from each SDF's coarse-texture dimensions. - -The helpers here exist solely so the deprecated ``Model.sdf_block_coords`` -and ``Model.sdf_index2blocks`` properties can keep returning equivalent -arrays for callers that still read them. They visit every subgrid in the -coarse grid (matching the new broadphase semantics), so the returned -coords are dense rather than narrow-band. -""" - -from __future__ import annotations - -import numpy as np -import warp as wp - - -def compute_block_coords_and_index2blocks( - coarse_textures: list, - subgrid_size: int = 8, -) -> tuple[np.ndarray, np.ndarray]: - """Build legacy ``(block_coords, index2blocks)`` arrays from coarse textures. - - For each coarse texture, emit every subgrid's voxel-space corner as a - ``vec3us``. The per-SDF range is recorded as ``[start, end)`` rows. - - Args: - coarse_textures: List of ``wp.Texture3D`` (one per SDF). Entries may - be ``None`` for SDFs that have no texture data. - subgrid_size: Cells per subgrid side. Matches the value baked into - ``TextureSDFData.subgrid_size`` at build time (default 8). - - Returns: - ``(block_coords, index2blocks)`` as NumPy arrays: - * ``block_coords`` — shape ``(N, 3)``, dtype ``uint16`` — - concatenated subgrid corner coordinates in voxel space. - * ``index2blocks`` — shape ``(num_sdfs, 2)``, dtype ``int32`` — - per-SDF ``[start, end)`` indices into ``block_coords``. - """ - coords_chunks: list[np.ndarray] = [] - index2blocks = np.zeros((len(coarse_textures), 2), dtype=np.int32) - cursor = 0 - - for sdf_idx, tex in enumerate(coarse_textures): - index2blocks[sdf_idx, 0] = cursor - if tex is not None: - cw = max(int(tex.width) - 1, 0) - ch = max(int(tex.height) - 1, 0) - cd = max(int(tex.depth) - 1, 0) - n = cw * ch * cd - if n > 0: - bz, by, bx = np.meshgrid( - np.arange(cd, dtype=np.uint16), - np.arange(ch, dtype=np.uint16), - np.arange(cw, dtype=np.uint16), - indexing="ij", - ) - sgs = np.uint16(subgrid_size) - chunk = np.stack( - [(bx * sgs).ravel(), (by * sgs).ravel(), (bz * sgs).ravel()], - axis=-1, - ).astype(np.uint16) - coords_chunks.append(chunk) - cursor += n - index2blocks[sdf_idx, 1] = cursor - - if coords_chunks: - block_coords = np.concatenate(coords_chunks, axis=0) - else: - block_coords = np.zeros((0, 3), dtype=np.uint16) - - return block_coords, index2blocks - - -def build_legacy_sdf_block_arrays( - coarse_textures: list, - subgrid_size: int = 8, - device: str | None = None, -) -> tuple[wp.array, wp.array]: - """Return the legacy ``(sdf_block_coords, sdf_index2blocks)`` Warp arrays. - - Thin wrapper around :func:`compute_block_coords_and_index2blocks` that - materializes the numpy results as Warp arrays with the historical - dtypes (``wp.vec3us`` and ``wp.vec2i``). - """ - block_coords_np, index2blocks_np = compute_block_coords_and_index2blocks(coarse_textures, subgrid_size=subgrid_size) - sdf_block_coords = wp.array(block_coords_np, dtype=wp.vec3us, device=device) - sdf_index2blocks = wp.array(index2blocks_np, dtype=wp.vec2i, device=device) - return sdf_block_coords, sdf_index2blocks diff --git a/newton/_src/geometry/bvh.py b/newton/_src/geometry/bvh.py index a580503d7b..f59908e41d 100644 --- a/newton/_src/geometry/bvh.py +++ b/newton/_src/geometry/bvh.py @@ -3,13 +3,11 @@ from __future__ import annotations -import warnings from typing import TYPE_CHECKING import warp as wp from ..core import MAXVAL -from .flags import ShapeFlags from .types import Gaussian, GeoType if TYPE_CHECKING: @@ -170,12 +168,13 @@ def is_supported_shape_type(shape_type: wp.int32) -> wp.bool: def compute_enabled_shapes( shape_type: wp.array[wp.int32], shape_flags: wp.array[wp.int32], + shape_flags_mask: wp.int32, out_shape_enabled: wp.array[wp.uint32], out_shape_enabled_count: wp.array[wp.int32], ): tid = wp.tid() - if not bool(shape_flags[tid] & ShapeFlags.VISIBLE): + if not bool(shape_flags[tid] & shape_flags_mask): return if not is_supported_shape_type(shape_type[tid]): @@ -393,44 +392,6 @@ def compute_shape_world_transforms_launch(model: Model, state: State) -> None: ) -def build_bvh_shape(model: Model, state: State, *, bvh_constructor: str | None = None) -> None: - """Deprecated alias for :meth:`newton.Model.bvh_build_shapes`. - - .. deprecated:: 1.3 - Use :meth:`newton.Model.bvh_build_shapes` instead. - - Args: - model: Simulation model providing shape metadata. - state: Current simulation state with body transforms. - bvh_constructor: Warp BVH construction algorithm forwarded to - :meth:`newton.Model.bvh_build_shapes`. - """ - warnings.warn( - "newton.geometry.build_bvh_shape(model, state) is deprecated; use model.bvh_build_shapes(state) instead.", - DeprecationWarning, - stacklevel=2, - ) - model.bvh_build_shapes(state, bvh_constructor=bvh_constructor) - - -def refit_bvh_shape(model: Model, state: State) -> None: - """Deprecated alias for :meth:`newton.Model.bvh_refit_shapes`. - - .. deprecated:: 1.3 - Use :meth:`newton.Model.bvh_refit_shapes` instead. - - Args: - model: Simulation model providing shape metadata. - state: Current simulation state with body transforms. - """ - warnings.warn( - "newton.geometry.refit_bvh_shape(model, state) is deprecated; use model.bvh_refit_shapes(state) instead.", - DeprecationWarning, - stacklevel=2, - ) - model.bvh_refit_shapes(state) - - def compute_particle_bvh_bounds_launch( model: Model, state: State, @@ -454,41 +415,3 @@ def compute_particle_bvh_bounds_launch( ], device=model.device, ) - - -def build_bvh_particle(model: Model, state: State, *, bvh_constructor: str | None = None) -> None: - """Deprecated alias for :meth:`newton.Model.bvh_build_particles`. - - .. deprecated:: 1.3 - Use :meth:`newton.Model.bvh_build_particles` instead. - - Args: - model: Simulation model providing particle metadata. - state: Current simulation state with particle positions. - bvh_constructor: Warp BVH construction algorithm forwarded to - :meth:`newton.Model.bvh_build_particles`. - """ - warnings.warn( - "newton.geometry.build_bvh_particle(model, state) is deprecated; use model.bvh_build_particles(state) instead.", - DeprecationWarning, - stacklevel=2, - ) - model.bvh_build_particles(state, bvh_constructor=bvh_constructor) - - -def refit_bvh_particle(model: Model, state: State) -> None: - """Deprecated alias for :meth:`newton.Model.bvh_refit_particles`. - - .. deprecated:: 1.3 - Use :meth:`newton.Model.bvh_refit_particles` instead. - - Args: - model: Simulation model providing particle metadata. - state: Current simulation state with particle positions. - """ - warnings.warn( - "newton.geometry.refit_bvh_particle(model, state) is deprecated; use model.bvh_refit_particles(state) instead.", - DeprecationWarning, - stacklevel=2, - ) - model.bvh_refit_particles(state) diff --git a/newton/_src/geometry/inertia.py b/newton/_src/geometry/inertia.py index 9b09dc58d1..4a8e4c3bf4 100644 --- a/newton/_src/geometry/inertia.py +++ b/newton/_src/geometry/inertia.py @@ -5,7 +5,9 @@ from __future__ import annotations +import math import warnings +from numbers import Real import numpy as np import warp as wp @@ -41,6 +43,36 @@ _MESH_INERTIA_TILE_SIZE = 256 +def _validate_hollow_thickness( + shape_name: str, + thickness: float, + limits: tuple[tuple[str, float], ...], +) -> float: + if isinstance(thickness, bool) or not isinstance(thickness, Real): + raise TypeError(f"thickness must be a real scalar for a hollow {shape_name} geom") + + thickness = float(thickness) + if not math.isfinite(thickness): + raise ValueError(f"thickness must be finite for a hollow {shape_name} geom; got {thickness}") + if thickness < 0.0: + raise ValueError(f"thickness must be >= 0 for a hollow {shape_name} geom; got {thickness}") + if thickness == 0.0: + warnings.warn( + f"A hollow {shape_name} geom with zero thickness has zero mass and inertia.", + stacklevel=2, + ) + return thickness + + for dim_name, dim_value in limits: + dim_limit = float(dim_value) + if thickness >= dim_limit: + raise ValueError( + f"thickness ({thickness}) must be smaller than {dim_name} ({dim_limit}) for a hollow {shape_name} geom" + ) + + return thickness + + def compute_inertia_sphere(density: float, radius: float) -> tuple[float, wp.vec3, wp.mat33]: """Helper to compute mass and inertia of a solid sphere @@ -582,7 +614,9 @@ def compute_inertia_shape( if is_solid: return solid else: - assert isinstance(thickness, float), "thickness must be a float for a hollow sphere geom" + thickness = _validate_hollow_thickness("sphere", thickness, (("radius", scale[0]),)) + if thickness == 0.0: + return 0.0, solid[1], wp.mat33() hollow = compute_inertia_sphere(density, scale[0] - thickness) return solid[0] - hollow[0], solid[1], solid[2] - hollow[2] elif type == GeoType.BOX: @@ -591,7 +625,11 @@ def compute_inertia_shape( if is_solid: return solid else: - assert isinstance(thickness, float), "thickness must be a float for a hollow box geom" + thickness = _validate_hollow_thickness( + "box", thickness, (("hx", scale[0]), ("hy", scale[1]), ("hz", scale[2])) + ) + if thickness == 0.0: + return 0.0, solid[1], wp.mat33() hollow = compute_inertia_box(density, scale[0] - thickness, scale[1] - thickness, scale[2] - thickness) return solid[0] - hollow[0], solid[1], solid[2] - hollow[2] elif type == GeoType.CAPSULE: @@ -600,7 +638,11 @@ def compute_inertia_shape( if is_solid: return solid else: - assert isinstance(thickness, float), "thickness must be a float for a hollow capsule geom" + thickness = _validate_hollow_thickness( + "capsule", thickness, (("radius", scale[0]), ("half_height", scale[1])) + ) + if thickness == 0.0: + return 0.0, solid[1], wp.mat33() hollow = compute_inertia_capsule(density, scale[0] - thickness, scale[1] - thickness) return solid[0] - hollow[0], solid[1], solid[2] - hollow[2] elif type == GeoType.CYLINDER: @@ -609,7 +651,11 @@ def compute_inertia_shape( if is_solid: return solid else: - assert isinstance(thickness, float), "thickness must be a float for a hollow cylinder geom" + thickness = _validate_hollow_thickness( + "cylinder", thickness, (("radius", scale[0]), ("half_height", scale[1])) + ) + if thickness == 0.0: + return 0.0, solid[1], wp.mat33() hollow = compute_inertia_cylinder(density, scale[0] - thickness, scale[1] - thickness) return solid[0] - hollow[0], solid[1], solid[2] - hollow[2] elif type == GeoType.CONE: @@ -618,15 +664,11 @@ def compute_inertia_shape( if is_solid: return solid else: - assert isinstance(thickness, float), "thickness must be a float for a hollow cone geom" + thickness = _validate_hollow_thickness("cone", thickness, (("radius", scale[0]), ("half_height", scale[1]))) + if thickness == 0.0: + return 0.0, solid[1], wp.mat33() hollow = compute_inertia_cone(density, scale[0] - thickness, scale[1] - thickness) m_shell = solid[0] - hollow[0] - if m_shell <= 0.0: - raise ValueError( - f"Hollow cone shell has non-positive mass ({m_shell:.6g}). " - f"The thickness ({thickness}) must be smaller than both the " - f"radius ({scale[0]}) and half_height ({scale[1]})." - ) # Cones have non-zero COM so outer and inner cones have different COMs; # compute the shell COM as the weighted difference, then shift both # inertia tensors to the shell COM before subtracting (parallel-axis theorem). @@ -648,7 +690,11 @@ def _shift_inertia(mass, I_mat, com_from, com_to): if is_solid: return solid else: - assert isinstance(thickness, float), "thickness must be a float for a hollow ellipsoid geom" + thickness = _validate_hollow_thickness( + "ellipsoid", thickness, (("rx", scale[0]), ("ry", scale[1]), ("rz", scale[2])) + ) + if thickness == 0.0: + return 0.0, solid[1], wp.mat33() hollow = compute_inertia_ellipsoid( density, scale[0] - thickness, scale[1] - thickness, scale[2] - thickness ) diff --git a/newton/_src/geometry/sdf_hydroelastic.py b/newton/_src/geometry/sdf_hydroelastic.py index 91d668aab7..790328ea88 100644 --- a/newton/_src/geometry/sdf_hydroelastic.py +++ b/newton/_src/geometry/sdf_hydroelastic.py @@ -171,7 +171,7 @@ def mc_calc_face_texture( p_0 = wp.vec3f(corner_offsets_table[v_idx_from]) p_1 = wp.vec3f(corner_offsets_table[v_idx_to]) val_diff = wp.float32(val_1 - val_0) - if wp.abs(val_diff) < wp.static(MC_EDGE_VAL_DIFF_EPS): + if wp.abs(val_diff) < MC_EDGE_VAL_DIFF_EPS: t = float(0.5) else: # Clamp t away from cube corners to prevent vertex collapse when @@ -194,7 +194,7 @@ def mc_calc_face_texture( n = wp.cross(face_verts[1] - face_verts[0], face_verts[2] - face_verts[0]) n_sq = wp.dot(n, n) - if n_sq < wp.static(MC_DEGENERATE_N_SQ_EPS): + if n_sq < MC_DEGENERATE_N_SQ_EPS: # Degenerate triangle — return zero area with a valid (non-NaN) normal. area = 0.0 normal = wp.vec3(0.0, 0.0, 1.0) @@ -580,16 +580,15 @@ def _from_model( """ shape_flags = model.shape_flags.numpy() - # Check if any shapes have hydroelastic flag - has_hydroelastic = any((flags & ShapeFlags.HYDROELASTIC) for flags in shape_flags) - if not has_hydroelastic: + # Check if any shapes have hydroelastic flag. + is_hydroelastic = (shape_flags & int(ShapeFlags.HYDROELASTIC)) != 0 + if not is_hydroelastic.any(): return None - shape_pairs = model.shape_contact_pairs.numpy() - num_hydroelastic_pairs = 0 - for shape_a, shape_b in shape_pairs: - if (shape_flags[shape_a] & ShapeFlags.HYDROELASTIC) and (shape_flags[shape_b] & ShapeFlags.HYDROELASTIC): - num_hydroelastic_pairs += 1 + shape_pairs = model.shape_contact_pairs.numpy().reshape(-1, 2) + num_hydroelastic_pairs = int( + np.count_nonzero(is_hydroelastic[shape_pairs[:, 0]] & is_hydroelastic[shape_pairs[:, 1]]) + ) if num_hydroelastic_pairs == 0: return None diff --git a/newton/_src/geometry/sdf_mc.py b/newton/_src/geometry/sdf_mc.py index 4e448df919..4b1bb4403b 100644 --- a/newton/_src/geometry/sdf_mc.py +++ b/newton/_src/geometry/sdf_mc.py @@ -195,14 +195,14 @@ def mc_calc_face( p_0 = wp.vec3f(corner_offsets_table[v_idx_from]) p_1 = wp.vec3f(corner_offsets_table[v_idx_to]) val_diff = wp.float32(val_1 - val_0) - if wp.abs(val_diff) < wp.static(MC_EDGE_VAL_DIFF_EPS): + if wp.abs(val_diff) < MC_EDGE_VAL_DIFF_EPS: p = 0.5 * (p_0 + p_1) else: # Clamp t away from cube corners to prevent vertex collapse when # corner values are near zero (e.g. at SDF ridge boundaries). # Without the clamp, t close to 0 or 1 places multiple vertices # at the same corner, producing degenerate (zero-area) triangles. - t = wp.clamp((isovalue - val_0) / val_diff, wp.static(MC_EDGE_CLAMP_MIN), wp.static(MC_EDGE_CLAMP_MAX)) + t = wp.clamp((isovalue - val_0) / val_diff, MC_EDGE_CLAMP_MIN, MC_EDGE_CLAMP_MAX) p = p_0 + t * (p_1 - p_0) vol_idx = p + int_to_vec3f(x_id, y_id, z_id) p_scaled = wp.volume_index_to_world(sdf_a, vol_idx) @@ -216,7 +216,7 @@ def mc_calc_face( n = wp.cross(face_verts[1] - face_verts[0], face_verts[2] - face_verts[0]) n_sq = wp.dot(n, n) - if n_sq < wp.static(MC_DEGENERATE_N_SQ_EPS): + if n_sq < MC_DEGENERATE_N_SQ_EPS: # Degenerate triangle — return zero area with a valid (non-NaN) normal. area = 0.0 normal = wp.vec3(0.0, 0.0, 1.0) diff --git a/newton/_src/geometry/sdf_texture.py b/newton/_src/geometry/sdf_texture.py index 6e01679d6c..e0f3708219 100644 --- a/newton/_src/geometry/sdf_texture.py +++ b/newton/_src/geometry/sdf_texture.py @@ -6,9 +6,9 @@ This module provides a GPU-accelerated sparse SDF implementation using 3D CUDA textures. Construction mirrors the NanoVDB sparse-volume pattern in ``sdf_utils.py``: -1. Check subgrid occupancy by querying mesh SDF at subgrid centers -2. Build background/coarse SDF by querying mesh at subgrid corner positions -3. Populate only occupied subgrid textures by querying mesh at each texel +1. Check subgrid occupancy by querying the source SDF at subgrid centers +2. Build the background/coarse SDF by querying the source at subgrid corner positions +3. Populate only occupied subgrid textures by querying the source at each texel The format uses: - A coarse 3D texture for background/far-field sampling @@ -23,17 +23,26 @@ import functools import types +from collections.abc import Callable, Sequence import numpy as np import warp as wp +from ..core.types import Axis +from .kernels import sdf_box, sdf_capsule, sdf_cone, sdf_cylinder, sdf_ellipsoid, sdf_sphere from .sdf_mc import MC_EDGE_CLAMP_MAX, MC_EDGE_CLAMP_MIN, MC_EDGE_VAL_DIFF_EPS -from .sdf_utils import get_distance_to_mesh, get_distance_to_mesh_normal, get_distance_to_mesh_parity +from .sdf_utils import ( + get_distance_to_mesh, + get_distance_to_mesh_normal, + get_distance_to_mesh_parity, + get_primitive_extents, +) +from .types import GeoType # Sentinel values for subgrid indirection slots. -# Plain int so wp.static() works in kernels; numpy casts on assignment. -SLOT_EMPTY = 0xFFFFFFFF # No subgrid data (empty/far-field cell) -SLOT_LINEAR = 0xFFFFFFFE # Subgrid demoted to coarse interpolation +# Typed uint32 so kernel codegen doesn't overflow an int32 constant. +SLOT_EMPTY = wp.uint32(0xFFFFFFFF) # No subgrid data (empty/far-field cell) +SLOT_LINEAR = wp.uint32(0xFFFFFFFE) # Subgrid demoted to coarse interpolation # Inside/outside sign strategies for the mesh distance queries during the # bake. Winding is 0 so a legacy ``use_parity`` boolean maps onto the same @@ -42,6 +51,39 @@ SIGN_MODE_PARITY = 1 SIGN_MODE_NORMAL = 2 +_SDF_SOURCE_MESH = 0 +_SDF_SOURCE_PRIMITIVE = 1 + +_SUPPORTED_PRIMITIVE_SDF_TYPES = frozenset( + { + int(GeoType.SPHERE), + int(GeoType.BOX), + int(GeoType.CAPSULE), + int(GeoType.CYLINDER), + int(GeoType.ELLIPSOID), + int(GeoType.CONE), + } +) + + +def _validate_primitive_sdf_inputs( + shape_type: int, + shape_scale: Sequence[float], +) -> tuple[int, tuple[float, float, float]]: + shape_type = int(shape_type) + if shape_type not in _SUPPORTED_PRIMITIVE_SDF_TYPES: + raise NotImplementedError(f"Texture SDF generation is not implemented for shape type: {shape_type}") + if len(shape_scale) != 3: + raise ValueError("shape_scale must contain exactly 3 components") + + scale = tuple(float(v) for v in shape_scale) + if not np.all(np.isfinite(scale)): + raise ValueError("shape_scale components must be finite") + if any(v < 0.0 for v in scale): + raise ValueError("shape_scale components must be non-negative") + return shape_type, scale + + # ============================================================================ # SDF texture-sampling paths # ============================================================================ @@ -137,6 +179,25 @@ def _id_to_xyz(idx: int, size_x: int, size_y: int) -> wp.vec3i: return wp.vec3i(x, y, z) +@wp.func +def _query_primitive_sdf(shape_type: wp.int32, shape_scale: wp.vec3, point: wp.vec3) -> float: + """Evaluate an analytical primitive SDF in the shape-local frame.""" + signed_distance = float(1.0e6) + if shape_type == GeoType.SPHERE: + signed_distance = sdf_sphere(point, shape_scale[0]) + elif shape_type == GeoType.BOX: + signed_distance = sdf_box(point, shape_scale[0], shape_scale[1], shape_scale[2]) + elif shape_type == GeoType.CAPSULE: + signed_distance = sdf_capsule(point, shape_scale[0], shape_scale[1], int(Axis.Z)) + elif shape_type == GeoType.CYLINDER: + signed_distance = sdf_cylinder(point, shape_scale[0], shape_scale[1], int(Axis.Z)) + elif shape_type == GeoType.ELLIPSOID: + signed_distance = sdf_ellipsoid(point, shape_scale) + elif shape_type == GeoType.CONE: + signed_distance = sdf_cone(point, shape_scale[0], shape_scale[1], int(Axis.Z)) + return signed_distance + + @wp.func def _interp_coarse_sdf( background_sdf: wp.array[float], @@ -217,43 +278,71 @@ def _write_subgrid_slot( @functools.cache -def _create_sign_mode_kernels(sign_mode: int): - """Generate the mesh-bake kernels specialized for one sign strategy. +def _create_source_kernels(source_kind: int, sign_mode: int): + """Generate bake kernels specialized for one source and sign strategy. The query dispatch is resolved at kernel-generation time instead of a - runtime branch: compiling all three query paths into every bake kernel - (in particular the heavy pseudo-normal query) raised register pressure - enough to slow ``build_sdf`` by ~19% on CI even when the extra branch - was never taken. One kernel set is generated and cached per sign mode; - the default parity/winding bakes never compile the pseudo-normal query. + runtime branch. This keeps analytic primitive queries separate from mesh + queries and ensures the default parity/winding mesh bakes never compile + the heavier pseudo-normal query. """ - if sign_mode == SIGN_MODE_PARITY: + if source_kind == _SDF_SOURCE_PRIMITIVE: @wp.func - def query_mesh_sdf( - mesh: wp.uint64, point: wp.vec3, max_dist: wp.float32, winding_threshold: wp.float32 + def query_sdf( + mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, + point: wp.vec3, + max_dist: wp.float32, + winding_threshold: wp.float32, + ) -> float: + return _query_primitive_sdf(shape_type, shape_scale, point) + + elif sign_mode == SIGN_MODE_PARITY: + + @wp.func + def query_sdf( + mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, + point: wp.vec3, + max_dist: wp.float32, + winding_threshold: wp.float32, ) -> float: return get_distance_to_mesh_parity(mesh, point, max_dist) elif sign_mode == SIGN_MODE_NORMAL: @wp.func - def query_mesh_sdf( - mesh: wp.uint64, point: wp.vec3, max_dist: wp.float32, winding_threshold: wp.float32 + def query_sdf( + mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, + point: wp.vec3, + max_dist: wp.float32, + winding_threshold: wp.float32, ) -> float: return get_distance_to_mesh_normal(mesh, point, max_dist) else: @wp.func - def query_mesh_sdf( - mesh: wp.uint64, point: wp.vec3, max_dist: wp.float32, winding_threshold: wp.float32 + def query_sdf( + mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, + point: wp.vec3, + max_dist: wp.float32, + winding_threshold: wp.float32, ) -> float: return get_distance_to_mesh(mesh, point, max_dist, winding_threshold) @wp.kernel def check_subgrid_occupied_kernel( mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, threshold: wp.vec2f, winding_threshold: float, subgrid_required: wp.array[wp.int32], @@ -263,7 +352,7 @@ def check_subgrid_occupied_kernel( min_corner: wp.vec3, cell_size: wp.vec3, ): - """Mark subgrids that overlap the narrow band by checking mesh SDF at center.""" + """Mark subgrids that overlap the narrow band by checking the source SDF at the center.""" tid = wp.tid() coords = _id_to_xyz(tid, num_subgrids_x, num_subgrids_y) sample_pos = min_corner + wp.vec3( @@ -272,7 +361,7 @@ def check_subgrid_occupied_kernel( (float(coords[2] * cells_per_subgrid) + float(cells_per_subgrid) * 0.5) * cell_size[2], ) - signed_distance = query_mesh_sdf(mesh, sample_pos, 10000.0, winding_threshold) + signed_distance = query_sdf(mesh, shape_type, shape_scale, sample_pos, 10000.0, winding_threshold) if _is_in_narrow_band(signed_distance, threshold): subgrid_required[tid] = 1 else: @@ -281,6 +370,8 @@ def check_subgrid_occupied_kernel( @wp.kernel def accumulate_subgrid_linearity_error_kernel( mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, background_sdf: wp.array[float], subgrid_required: wp.array[wp.int32], linearity_errors: wp.array[float], @@ -295,7 +386,7 @@ def accumulate_subgrid_linearity_error_kernel( bg_size_y: int, bg_size_z: int, ): - """Sample mesh SDF at every fine-grid point of every occupied subgrid and + """Sample the source SDF at every fine-grid point of every occupied subgrid and accumulate the maximum absolute deviation from the trilinearly interpolated coarse SDF via ``wp.atomic_max``. @@ -337,7 +428,7 @@ def accumulate_subgrid_linearity_error_kernel( float(gy) * cell_size[1], float(gz) * cell_size[2], ) - mesh_val = query_mesh_sdf(mesh, pos, 10000.0, winding_threshold) + sdf_value = query_sdf(mesh, shape_type, shape_scale, pos, 10000.0, winding_threshold) inv_cpsg = 1.0 / float(cells_per_subgrid) coarse_val = _interp_coarse_sdf( @@ -354,11 +445,13 @@ def accumulate_subgrid_linearity_error_kernel( bg_size_z, ) - wp.atomic_max(linearity_errors, subgrid_idx, wp.abs(mesh_val - coarse_val)) + wp.atomic_max(linearity_errors, subgrid_idx, wp.abs(sdf_value - coarse_val)) @wp.kernel - def build_coarse_sdf_from_mesh_kernel( + def build_coarse_sdf_kernel( mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, background_sdf: wp.array[float], min_corner: wp.vec3, cell_size: wp.vec3, @@ -368,7 +461,7 @@ def build_coarse_sdf_from_mesh_kernel( bg_size_z: int, winding_threshold: float, ): - """Populate background SDF by querying mesh at subgrid corner positions.""" + """Populate the background SDF by querying the source at subgrid corner positions.""" tid = wp.tid() total_bg = bg_size_x * bg_size_y * bg_size_z @@ -386,11 +479,13 @@ def build_coarse_sdf_from_mesh_kernel( float(z_block * cells_per_subgrid) * cell_size[2], ) - background_sdf[tid] = query_mesh_sdf(mesh, pos, 10000.0, winding_threshold) + background_sdf[tid] = query_sdf(mesh, shape_type, shape_scale, pos, 10000.0, winding_threshold) @wp.kernel def populate_subgrid_texture_float32_kernel( mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, subgrid_required: wp.array[wp.int32], subgrid_addresses: wp.array[wp.int32], subgrid_start_slots: wp.array3d[wp.uint32], @@ -405,7 +500,7 @@ def populate_subgrid_texture_float32_kernel( tex_blocks_per_dim: int, tex_size: int, ): - """Populate subgrid texture by querying mesh SDF (float32 version).""" + """Populate the subgrid texture by querying the source SDF (float32 version).""" tid = wp.tid() total_subgrids = num_subgrids_x * num_subgrids_y * num_subgrids_z @@ -439,7 +534,7 @@ def populate_subgrid_texture_float32_kernel( float(gy) * cell_size[1], float(gz) * cell_size[2], ) - sdf_val = query_mesh_sdf(mesh, pos, 10000.0, winding_threshold) + sdf_val = query_sdf(mesh, shape_type, shape_scale, pos, 10000.0, winding_threshold) address = subgrid_addresses[subgrid_idx] if address < 0: @@ -456,6 +551,8 @@ def populate_subgrid_texture_float32_kernel( @wp.kernel def populate_subgrid_texture_uint16_kernel( mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, subgrid_required: wp.array[wp.int32], subgrid_addresses: wp.array[wp.int32], subgrid_start_slots: wp.array3d[wp.uint32], @@ -472,7 +569,7 @@ def populate_subgrid_texture_uint16_kernel( sdf_min: float, sdf_range_inv: float, ): - """Populate subgrid texture by querying mesh SDF (uint16 quantized version).""" + """Populate the subgrid texture by querying the source SDF (uint16 quantized version).""" tid = wp.tid() total_subgrids = num_subgrids_x * num_subgrids_y * num_subgrids_z @@ -506,7 +603,7 @@ def populate_subgrid_texture_uint16_kernel( float(gy) * cell_size[1], float(gz) * cell_size[2], ) - sdf_val = query_mesh_sdf(mesh, pos, 10000.0, winding_threshold) + sdf_val = query_sdf(mesh, shape_type, shape_scale, pos, 10000.0, winding_threshold) address = subgrid_addresses[subgrid_idx] if address < 0: @@ -524,6 +621,8 @@ def populate_subgrid_texture_uint16_kernel( @wp.kernel def populate_subgrid_texture_uint8_kernel( mesh: wp.uint64, + shape_type: wp.int32, + shape_scale: wp.vec3, subgrid_required: wp.array[wp.int32], subgrid_addresses: wp.array[wp.int32], subgrid_start_slots: wp.array3d[wp.uint32], @@ -540,7 +639,7 @@ def populate_subgrid_texture_uint8_kernel( sdf_min: float, sdf_range_inv: float, ): - """Populate subgrid texture by querying mesh SDF (uint8 quantized version).""" + """Populate the subgrid texture by querying the source SDF (uint8 quantized version).""" tid = wp.tid() total_subgrids = num_subgrids_x * num_subgrids_y * num_subgrids_z @@ -574,7 +673,7 @@ def populate_subgrid_texture_uint8_kernel( float(gy) * cell_size[1], float(gz) * cell_size[2], ) - sdf_val = query_mesh_sdf(mesh, pos, 10000.0, winding_threshold) + sdf_val = query_sdf(mesh, shape_type, shape_scale, pos, 10000.0, winding_threshold) address = subgrid_addresses[subgrid_idx] if address < 0: @@ -596,7 +695,7 @@ def populate_subgrid_texture_uint8_kernel( return types.SimpleNamespace( check_subgrid_occupied_kernel=check_subgrid_occupied_kernel, accumulate_subgrid_linearity_error_kernel=accumulate_subgrid_linearity_error_kernel, - build_coarse_sdf_from_mesh_kernel=build_coarse_sdf_from_mesh_kernel, + build_coarse_sdf_kernel=build_coarse_sdf_kernel, populate_subgrid_texture_float32_kernel=populate_subgrid_texture_float32_kernel, populate_subgrid_texture_uint16_kernel=populate_subgrid_texture_uint16_kernel, populate_subgrid_texture_uint8_kernel=populate_subgrid_texture_uint8_kernel, @@ -757,7 +856,7 @@ def _read_cell_corners( ty = loc.ty tz = loc.tz - if loc.start_slot >= wp.static(SLOT_LINEAR): + if loc.start_slot >= SLOT_LINEAR: cx = float(loc.x_base) cy = float(loc.y_base) cz = float(loc.z_base) @@ -852,7 +951,7 @@ def texture_sample_sdf_at_voxel( start_slot = sdf.subgrid_start_slots[x_base, y_base, z_base] - if start_slot < wp.static(SLOT_LINEAR): + if start_slot < SLOT_LINEAR: block_x = float(start_slot & wp.uint32(0x3FF)) block_y = float((start_slot >> wp.uint32(10)) & wp.uint32(0x3FF)) block_z = float((start_slot >> wp.uint32(20)) & wp.uint32(0x3FF)) @@ -924,7 +1023,7 @@ def texture_sample_sdf( ty = loc.ty tz = loc.tz - if loc.start_slot >= wp.static(SLOT_LINEAR): + if loc.start_slot >= SLOT_LINEAR: cx = float(loc.x_base) cy = float(loc.y_base) cz = float(loc.z_base) @@ -1008,7 +1107,7 @@ def texture_sample_sdf_hw( sdf_val = float(0.0) - if loc.start_slot >= wp.static(SLOT_LINEAR): + if loc.start_slot >= SLOT_LINEAR: # ``cx + tx + 0.5`` lands at the centre of voxel (cx, cy, cz) and # ``+tx`` walks toward (cx+1, ...). The HW filter returns the # interpolated value in one fetch. @@ -1209,8 +1308,11 @@ def texture_sample_sdf_grad_only_hw( # ============================================================================ -def build_sparse_sdf_from_mesh( - mesh: wp.Mesh, +def _build_sparse_sdf( + source_kind: int, + mesh: wp.Mesh | None, + shape_type: int, + shape_scale: Sequence[float], grid_size_x: int, grid_size_y: int, grid_size_z: int, @@ -1225,7 +1327,7 @@ def build_sparse_sdf_from_mesh( sign_mode: int = SIGN_MODE_WINDING, device: str = "cuda", ) -> dict: - """Build sparse SDF texture representation by querying mesh directly. + """Build sparse SDF texture representation by querying a source directly. Mirrors the NanoVDB sparse-volume construction pattern: check subgrid occupancy at centers, then populate only occupied subgrids. Linearity @@ -1234,8 +1336,10 @@ def build_sparse_sdf_from_mesh( texture memory. Args: - mesh: Warp mesh. Must have ``support_winding_number=True`` unless - *sign_mode* is :data:`SIGN_MODE_PARITY` or :data:`SIGN_MODE_NORMAL`. + source_kind: SDF query source, either mesh or analytic primitive. + mesh: Warp mesh for mesh query sources. + shape_type: Primitive :class:`GeoType` for primitive query sources. + shape_scale: Primitive shape scale [m]. grid_size_x: fine grid X dimension [sample]. grid_size_y: fine grid Y dimension [sample]. grid_size_z: fine grid Z dimension [sample]. @@ -1276,6 +1380,9 @@ def build_sparse_sdf_from_mesh( min_corner_wp = wp.vec3(float(min_corner[0]), float(min_corner[1]), float(min_corner[2])) cell_size_wp = wp.vec3(float(cell_size[0]), float(cell_size[1]), float(cell_size[2])) + mesh_id = mesh.id if mesh is not None else wp.uint64(0) + shape_type_wp = wp.int32(shape_type) + shape_scale_wp = wp.vec3(float(shape_scale[0]), float(shape_scale[1]), float(shape_scale[2])) bg_size_x = w + 1 bg_size_y = h + 1 @@ -1292,14 +1399,16 @@ def build_sparse_sdf_from_mesh( # texture is sized so that subgrids whose SDF is well-approximated by # the coarse grid consume no high-resolution texture memory. # ------------------------------------------------------------------- - sign_kernels = _create_sign_mode_kernels(int(sign_mode)) + source_kernels = _create_source_kernels(int(source_kind), int(sign_mode)) background_sdf = wp.zeros(total_bg, dtype=float, device=device) wp.launch( - sign_kernels.build_coarse_sdf_from_mesh_kernel, + source_kernels.build_coarse_sdf_kernel, dim=total_bg, inputs=[ - mesh.id, + mesh_id, + shape_type_wp, + shape_scale_wp, background_sdf, min_corner_wp, cell_size_wp, @@ -1315,10 +1424,12 @@ def build_sparse_sdf_from_mesh( subgrid_required = wp.zeros(total_subgrids, dtype=wp.int32, device=device) threshold = wp.vec2f(-narrow_band_thickness - subgrid_radius, narrow_band_thickness + subgrid_radius) wp.launch( - sign_kernels.check_subgrid_occupied_kernel, + source_kernels.check_subgrid_occupied_kernel, dim=total_subgrids, inputs=[ - mesh.id, + mesh_id, + shape_type_wp, + shape_scale_wp, threshold, winding_threshold, subgrid_required, @@ -1338,10 +1449,10 @@ def build_sparse_sdf_from_mesh( if linearization_error_threshold > 0.0: # Per-sample launch so the 9^3 inner loop is parallelized across # threads; atomic_max accumulates the per-subgrid linearity error. - # We deliberately do NOT cache the mesh samples for reuse in the + # We deliberately do NOT cache the source samples for reuse in the # populate pass: an empirical test showed the cache (one float32 # per sample, total_subgrids * 9^3 bytes transient) costs more in - # global-memory traffic than re-querying the mesh BVH, both for + # global-memory traffic than re-querying the source SDF. This was measured for # small meshes (cube: 12 tris) and medium meshes (icosphere: # 5120 tris) at resolutions up to 256. samples_per_dim = subgrid_size + 1 @@ -1350,10 +1461,12 @@ def build_sparse_sdf_from_mesh( linearity_errors = wp.zeros(total_subgrids, dtype=float, device=device) wp.launch( - sign_kernels.accumulate_subgrid_linearity_error_kernel, + source_kernels.accumulate_subgrid_linearity_error_kernel, dim=total_work, inputs=[ - mesh.id, + mesh_id, + shape_type_wp, + shape_scale_wp, background_sdf, subgrid_required, linearity_errors, @@ -1426,10 +1539,12 @@ def build_sparse_sdf_from_mesh( if quantization_mode == QuantizationMode.FLOAT32: subgrid_texture_gpu = wp.zeros(total_tex_samples, dtype=float, device=device) wp.launch( - sign_kernels.populate_subgrid_texture_float32_kernel, + source_kernels.populate_subgrid_texture_float32_kernel, dim=total_work, inputs=[ - mesh.id, + mesh_id, + shape_type_wp, + shape_scale_wp, subgrid_required, subgrid_addresses, subgrid_start_slots_gpu, @@ -1452,10 +1567,12 @@ def build_sparse_sdf_from_mesh( elif quantization_mode == QuantizationMode.UINT16: subgrid_texture_gpu = wp.zeros(total_tex_samples, dtype=wp.uint16, device=device) wp.launch( - sign_kernels.populate_subgrid_texture_uint16_kernel, + source_kernels.populate_subgrid_texture_uint16_kernel, dim=total_work, inputs=[ - mesh.id, + mesh_id, + shape_type_wp, + shape_scale_wp, subgrid_required, subgrid_addresses, subgrid_start_slots_gpu, @@ -1480,10 +1597,12 @@ def build_sparse_sdf_from_mesh( elif quantization_mode == QuantizationMode.UINT8: subgrid_texture_gpu = wp.zeros(total_tex_samples, dtype=wp.uint8, device=device) wp.launch( - sign_kernels.populate_subgrid_texture_uint8_kernel, + source_kernels.populate_subgrid_texture_uint8_kernel, dim=total_work, inputs=[ - mesh.id, + mesh_id, + shape_type_wp, + shape_scale_wp, subgrid_required, subgrid_addresses, subgrid_start_slots_gpu, @@ -1548,6 +1667,122 @@ def build_sparse_sdf_from_mesh( } +def build_sparse_sdf_from_mesh( + mesh: wp.Mesh, + grid_size_x: int, + grid_size_y: int, + grid_size_z: int, + cell_size: np.ndarray, + min_corner: np.ndarray, + max_corner: np.ndarray, + subgrid_size: int = 8, + narrow_band_thickness: float = 0.1, + quantization_mode: int = QuantizationMode.UINT16, + winding_threshold: float = 0.5, + linearization_error_threshold: float | None = None, + sign_mode: int = SIGN_MODE_WINDING, + device: str = "cuda", +) -> dict: + """Build sparse SDF texture representation by querying a mesh directly. + + Args: + mesh: Warp mesh. Must have ``support_winding_number=True`` unless + *sign_mode* is :data:`SIGN_MODE_PARITY` or :data:`SIGN_MODE_NORMAL`. + grid_size_x: fine grid X dimension [sample]. + grid_size_y: fine grid Y dimension [sample]. + grid_size_z: fine grid Z dimension [sample]. + cell_size: fine grid cell size per axis [m]. + min_corner: lower corner of domain [m]. + max_corner: upper corner of domain [m]. + subgrid_size: cells per subgrid. + narrow_band_thickness: distance threshold for subgrids [m]. + quantization_mode: :class:`QuantizationMode` value. + winding_threshold: winding number threshold for inside/outside. + linearization_error_threshold: maximum absolute SDF error [m] below + which an occupied subgrid is considered linear. + sign_mode: inside/outside strategy for mesh distance queries. + device: Warp device string. + + Returns: + Dictionary with all sparse SDF data. + """ + return _build_sparse_sdf( + _SDF_SOURCE_MESH, + mesh, + 0, + (1.0, 1.0, 1.0), + grid_size_x, + grid_size_y, + grid_size_z, + cell_size, + min_corner, + max_corner, + subgrid_size=subgrid_size, + narrow_band_thickness=narrow_band_thickness, + quantization_mode=quantization_mode, + winding_threshold=winding_threshold, + linearization_error_threshold=linearization_error_threshold, + sign_mode=sign_mode, + device=device, + ) + + +def build_sparse_sdf_from_primitive( + shape_type: int, + shape_scale: Sequence[float], + grid_size_x: int, + grid_size_y: int, + grid_size_z: int, + cell_size: np.ndarray, + min_corner: np.ndarray, + max_corner: np.ndarray, + subgrid_size: int = 8, + narrow_band_thickness: float = 0.1, + quantization_mode: int = QuantizationMode.UINT16, + linearization_error_threshold: float | None = None, + device: str = "cuda", +) -> dict: + """Build sparse SDF texture representation from an analytical primitive. + + Args: + shape_type: Primitive :class:`GeoType`. + shape_scale: Primitive shape scale [m]. + grid_size_x: fine grid X dimension [sample]. + grid_size_y: fine grid Y dimension [sample]. + grid_size_z: fine grid Z dimension [sample]. + cell_size: fine grid cell size per axis [m]. + min_corner: lower corner of domain [m]. + max_corner: upper corner of domain [m]. + subgrid_size: cells per subgrid. + narrow_band_thickness: distance threshold for subgrids [m]. + quantization_mode: :class:`QuantizationMode` value. + linearization_error_threshold: maximum absolute SDF error [m] below + which an occupied subgrid is considered linear. + device: Warp device string. + + Returns: + Dictionary with all sparse SDF data. + """ + shape_type, shape_scale = _validate_primitive_sdf_inputs(shape_type, shape_scale) + return _build_sparse_sdf( + _SDF_SOURCE_PRIMITIVE, + None, + shape_type, + shape_scale, + grid_size_x, + grid_size_y, + grid_size_z, + cell_size, + min_corner, + max_corner, + subgrid_size=subgrid_size, + narrow_band_thickness=narrow_band_thickness, + quantization_mode=quantization_mode, + linearization_error_threshold=linearization_error_threshold, + device=device, + ) + + def create_sparse_sdf_textures( sparse_data: dict, device: str = "cuda", @@ -1555,7 +1790,7 @@ def create_sparse_sdf_textures( """Create TextureSDFData struct with GPU textures from sparse data. Args: - sparse_data: dictionary from :func:`build_sparse_sdf_from_mesh`. + sparse_data: Dictionary produced by a sparse SDF builder. device: Warp device string. Returns: @@ -1616,6 +1851,69 @@ def create_sparse_sdf_textures( return sdf_params, coarse_tex, subgrid_tex +def _create_texture_sdf_from_source( + min_ext: np.ndarray, + max_ext: np.ndarray, + *, + sparse_sdf_builder: Callable[..., dict], + narrow_band_range: tuple[float, float], + max_resolution: int | None, + target_voxel_size: float | None, + subgrid_size: int, + quantization_mode: int, + scale_baked: bool, + device: str, + return_sparse_data: bool = False, +) -> tuple[TextureSDFData, wp.Texture3D, wp.Texture3D] | tuple[TextureSDFData, wp.Texture3D, wp.Texture3D, dict | None]: + """Create a texture SDF from source extents and a bound sparse-data builder.""" + ext = max_ext - min_ext + max_ext_scalar = np.max(ext) + if max_ext_scalar < 1e-10: + empty = (create_empty_texture_sdf_data(), None, None) + return (*empty, None) if return_sparse_data else empty + + if target_voxel_size is not None: + if target_voxel_size <= 0.0: + raise ValueError("target_voxel_size must be > 0") + derived_res = int(np.ceil(max_ext_scalar / float(target_voxel_size))) + derived_res = max(8, ((derived_res + 7) // 8) * 8) + max_resolution = derived_res + elif max_resolution is None: + max_resolution = 64 + + max_resolution = int(max_resolution) + if max_resolution <= 0: + raise ValueError("max_resolution must be > 0") + if max_resolution >= (1 << 16): + raise ValueError(f"max_resolution must be less than {1 << 16}") + + cell_size_scalar = max_ext_scalar / max_resolution + dims = np.ceil(ext / cell_size_scalar).astype(int) + 1 + grid_x, grid_y, grid_z = int(dims[0]), int(dims[1]), int(dims[2]) + cell_size = ext / (dims - 1) + narrow_band_thickness = max(abs(narrow_band_range[0]), abs(narrow_band_range[1])) + + sparse_data = sparse_sdf_builder( + grid_x, + grid_y, + grid_z, + cell_size, + min_ext, + max_ext, + subgrid_size=subgrid_size, + narrow_band_thickness=narrow_band_thickness, + quantization_mode=quantization_mode, + device=device, + ) + + sdf_params, coarse_tex, subgrid_tex = create_sparse_sdf_textures(sparse_data, device) + sdf_params.scale_baked = scale_baked + + if return_sparse_data: + return sdf_params, coarse_tex, subgrid_tex, sparse_data + return sdf_params, coarse_tex, subgrid_tex + + def create_texture_sdf_from_mesh( mesh: wp.Mesh, *, @@ -1683,62 +1981,85 @@ def create_texture_sdf_from_mesh( min_ext = mesh_min - margin max_ext = mesh_max + margin - # Compute grid dimensions (same math as the former build_dense_sdf) - ext = max_ext - min_ext - max_ext_scalar = np.max(ext) - if max_ext_scalar < 1e-10: - empty = (create_empty_texture_sdf_data(), None, None) - return (*empty, None) if return_sparse_data else empty + sparse_sdf_builder = functools.partial( + build_sparse_sdf_from_mesh, + mesh, + winding_threshold=winding_threshold, + sign_mode=sign_mode, + ) + return _create_texture_sdf_from_source( + min_ext, + max_ext, + sparse_sdf_builder=sparse_sdf_builder, + narrow_band_range=narrow_band_range, + max_resolution=max_resolution, + target_voxel_size=target_voxel_size, + subgrid_size=subgrid_size, + quantization_mode=quantization_mode, + scale_baked=scale_baked, + device=device, + return_sparse_data=return_sparse_data, + ) - # Resolve max_resolution, honoring target_voxel_size when provided. - # Mirrors the sparse SDF path in sdf_utils._compute_sdf_from_shape_impl - # so texture and sparse grids agree on resolution. - if target_voxel_size is not None: - if target_voxel_size <= 0.0: - raise ValueError("target_voxel_size must be > 0") - derived_res = int(np.ceil(max_ext_scalar / float(target_voxel_size))) - # Keep alignment with tiled SDF builders that operate on 8-voxel chunks. - derived_res = max(8, ((derived_res + 7) // 8) * 8) - max_resolution = derived_res - elif max_resolution is None: - max_resolution = 64 - max_resolution = int(max_resolution) - if max_resolution <= 0: - raise ValueError("max_resolution must be > 0") - if max_resolution >= (1 << 16): - raise ValueError(f"max_resolution must be less than {1 << 16}") +def create_texture_sdf_from_primitive( + shape_type: int, + shape_scale: Sequence[float], + *, + margin: float = 0.05, + narrow_band_range: tuple[float, float] = (-0.1, 0.1), + max_resolution: int | None = None, + target_voxel_size: float | None = None, + subgrid_size: int = 8, + quantization_mode: int = QuantizationMode.UINT16, + scale_baked: bool = False, + device: str = "cuda", +) -> tuple[TextureSDFData, wp.Texture3D, wp.Texture3D]: + """Create texture SDF from an analytical primitive. - cell_size_scalar = max_ext_scalar / max_resolution - dims = np.ceil(ext / cell_size_scalar).astype(int) + 1 - grid_x, grid_y, grid_z = int(dims[0]), int(dims[1]), int(dims[2]) - cell_size = ext / (dims - 1) + Args: + shape_type: Primitive :class:`GeoType`. + shape_scale: Primitive shape scale [m], shape [3]. + margin: extra AABB padding [m]. + narrow_band_range: signed narrow-band distance range [m] as ``(inner, outer)``. + max_resolution: maximum grid dimension [voxel]. Used when + ``target_voxel_size`` is not provided. Defaults to 64 when both + ``max_resolution`` and ``target_voxel_size`` are ``None``. + target_voxel_size: target voxel size [m] along the longest padded-AABB + axis. When provided, takes precedence over ``max_resolution``. + subgrid_size: cells per subgrid. + quantization_mode: :class:`QuantizationMode` value. + scale_baked: whether shape scale was baked into the SDF values. + device: Warp device string. - narrow_band_thickness = max(abs(narrow_band_range[0]), abs(narrow_band_range[1])) + Returns: + Tuple of ``(texture_sdf, coarse_texture, subgrid_texture)``. + Caller must keep texture references alive to prevent GC. + """ + shape_type, shape_scale = _validate_primitive_sdf_inputs(shape_type, shape_scale) - sparse_data = build_sparse_sdf_from_mesh( - mesh, - grid_x, - grid_y, - grid_z, - cell_size, + min_prim, max_prim = get_primitive_extents(shape_type, shape_scale) + min_ext = np.asarray(min_prim, dtype=float) - margin + max_ext = np.asarray(max_prim, dtype=float) + margin + + sparse_sdf_builder = functools.partial( + build_sparse_sdf_from_primitive, + shape_type, + shape_scale, + ) + return _create_texture_sdf_from_source( min_ext, max_ext, + sparse_sdf_builder=sparse_sdf_builder, + narrow_band_range=narrow_band_range, + max_resolution=max_resolution, + target_voxel_size=target_voxel_size, subgrid_size=subgrid_size, - narrow_band_thickness=narrow_band_thickness, quantization_mode=quantization_mode, - winding_threshold=winding_threshold, - sign_mode=sign_mode, + scale_baked=scale_baked, device=device, ) - sdf_params, coarse_tex, subgrid_tex = create_sparse_sdf_textures(sparse_data, device) - sdf_params.scale_baked = scale_baked - - if return_sparse_data: - return sdf_params, coarse_tex, subgrid_tex, sparse_data - return sdf_params, coarse_tex, subgrid_tex - def create_texture_sdf_from_volume( sparse_volume: wp.Volume, @@ -2205,10 +2526,10 @@ def _generate_isomesh_texture_kernel( p_0 = wp.vec3f(corner_offsets_table[v_from]) p_1 = wp.vec3f(corner_offsets_table[v_to]) val_diff = val_1 - val_0 - if wp.abs(val_diff) < wp.static(MC_EDGE_VAL_DIFF_EPS): + if wp.abs(val_diff) < MC_EDGE_VAL_DIFF_EPS: p = 0.5 * (p_0 + p_1) else: - t = wp.clamp((isovalue - val_0) / val_diff, wp.static(MC_EDGE_CLAMP_MIN), wp.static(MC_EDGE_CLAMP_MAX)) + t = wp.clamp((isovalue - val_0) / val_diff, MC_EDGE_CLAMP_MIN, MC_EDGE_CLAMP_MAX) p = p_0 + t * (p_1 - p_0) vol_idx = p + wp.vec3(float(x_id), float(y_id), float(z_id)) local_pos = sdf.sdf_box_lower + wp.cw_mul(vol_idx, sdf.voxel_size) diff --git a/newton/_src/geometry/sdf_utils.py b/newton/_src/geometry/sdf_utils.py index fe7f91ec21..f808445b5e 100644 --- a/newton/_src/geometry/sdf_utils.py +++ b/newton/_src/geometry/sdf_utils.py @@ -3,7 +3,6 @@ import logging import os -import warnings from collections.abc import Sequence from typing import TYPE_CHECKING, Literal @@ -213,28 +212,6 @@ def __init__( self._coarse_texture = _coarse_texture self._subgrid_texture = _subgrid_texture - @property - def texture_block_coords(self) -> None: - """Deprecated. Always returns ``None``. - - Texture-SDF block coordinates were removed when the hydroelastic - broadphase started deriving them arithmetically from the per-shape - coarse-texture dimensions. The attribute is retained for one - release cycle so existing callers do not break. - - .. deprecated:: 1.3 - This attribute will be removed in a future release. - """ - warnings.warn( - "SDF.texture_block_coords is deprecated and always returns None; " - "it will be removed in a future release. The hydroelastic broadphase " - "now derives block coordinates arithmetically from each SDF's " - "coarse-texture dimensions and no longer needs this attribute.", - DeprecationWarning, - stacklevel=2, - ) - return None - def to_kernel_data(self) -> SDFData: """Return kernel-facing SDF payload.""" return self.data @@ -1544,10 +1521,10 @@ def _generate_dense_mc_kernel( p_0 = wp.vec3f(corner_offsets_table[ev[0]]) p_1 = wp.vec3f(corner_offsets_table[ev[1]]) val_diff = val_1 - val_0 - if wp.abs(val_diff) < wp.static(MC_EDGE_VAL_DIFF_EPS): + if wp.abs(val_diff) < MC_EDGE_VAL_DIFF_EPS: p = 0.5 * (p_0 + p_1) else: - t = wp.clamp((0.0 - val_0) / val_diff, wp.static(MC_EDGE_CLAMP_MIN), wp.static(MC_EDGE_CLAMP_MAX)) + t = wp.clamp((0.0 - val_0) / val_diff, MC_EDGE_CLAMP_MIN, MC_EDGE_CLAMP_MAX) p = p_0 + t * (p_1 - p_0) local = base + p face_verts[vi] = wp.vec3( @@ -1557,7 +1534,7 @@ def _generate_dense_mc_kernel( ) n = wp.cross(face_verts[1] - face_verts[0], face_verts[2] - face_verts[0]) n_sq = wp.dot(n, n) - if n_sq < wp.static(MC_DEGENERATE_N_SQ_EPS): + if n_sq < MC_DEGENERATE_N_SQ_EPS: normal = wp.vec3(0.0, 0.0, 1.0) else: normal = n / wp.sqrt(n_sq) diff --git a/newton/_src/geometry/terrain_generator.py b/newton/_src/geometry/terrain_generator.py index 3023766796..9987074639 100644 --- a/newton/_src/geometry/terrain_generator.py +++ b/newton/_src/geometry/terrain_generator.py @@ -634,77 +634,95 @@ def create_mesh_heightfield( if extent_y <= 0: raise ValueError(f"extent_y must be positive, got {extent_y}") - # Create grid coordinates + # Vertex and index buffers are allocated once and filled in place. The intermediate + # column_stack/vstack chain this replaces built the index buffer in int64 (twice the + # width it is finally stored in) and copied every triangle several times, which + # dominates terrain construction at Isaac Lab grid sizes (millions of triangles). + n_grid = grid_size_x * grid_size_y + n_quad = (grid_size_x - 1) * (grid_size_y - 1) + n_side = 2 * (grid_size_x - 1) + 2 * (grid_size_y - 1) + + # Vertices: top surface followed by the bottom surface, both on the same XY grid. x = np.linspace(-extent_x / 2, extent_x / 2, grid_size_x) + center_x y = np.linspace(-extent_y / 2, extent_y / 2, grid_size_y) + center_y - X, Y = np.meshgrid(x, y, indexing="ij") - - # Top and bottom surface vertices - top_vertices = np.column_stack([X.ravel(), Y.ravel(), heightfield.ravel()]).astype(np.float32) - bottom_z = np.full_like(heightfield, ground_z) - bottom_vertices = np.column_stack([X.ravel(), Y.ravel(), bottom_z.ravel()]).astype(np.float32) - vertices = np.vstack([top_vertices, bottom_vertices]) - - # Generate quad indices for all grid cells - i_indices = np.arange(grid_size_x - 1) - j_indices = np.arange(grid_size_y - 1) - ii, jj = np.meshgrid(i_indices, j_indices, indexing="ij") - ii, jj = ii.ravel(), jj.ravel() - - v0 = ii * grid_size_y + jj - v1 = ii * grid_size_y + (jj + 1) - v2 = (ii + 1) * grid_size_y + jj - v3 = (ii + 1) * grid_size_y + (jj + 1) - - # Top surface faces (counter-clockwise) - top_faces = np.column_stack([np.column_stack([v0, v2, v1]), np.column_stack([v1, v2, v3])]).reshape(-1, 3) - - # Bottom surface faces (clockwise) - num_top_vertices = len(top_vertices) - bottom_faces = np.column_stack( - [ - np.column_stack([num_top_vertices + v0, num_top_vertices + v1, num_top_vertices + v2]), - np.column_stack([num_top_vertices + v1, num_top_vertices + v3, num_top_vertices + v2]), - ] - ).reshape(-1, 3) - - # Side wall faces (4 edges) - side_faces_list = [] - i_edge = np.arange(grid_size_x - 1) - j_edge = np.arange(grid_size_y - 1) - + vertices = np.empty((2 * n_grid, 3), dtype=np.float32) + top_xyz = vertices[:n_grid].reshape(grid_size_x, grid_size_y, 3) + bottom_xyz = vertices[n_grid:].reshape(grid_size_x, grid_size_y, 3) + top_xyz[..., 0] = x[:, None] + top_xyz[..., 1] = y[None, :] + top_xyz[..., 2] = heightfield + bottom_xyz[..., 0] = x[:, None] + bottom_xyz[..., 1] = y[None, :] + bottom_xyz[..., 2] = ground_z + + # Corner indices of every grid cell, in the same row-major order as the vertices. + num_top_vertices = n_grid + ii = np.arange(grid_size_x - 1, dtype=np.int32)[:, None] + jj = np.arange(grid_size_y - 1, dtype=np.int32)[None, :] + v0 = (ii * grid_size_y + jj).ravel() + v1 = v0 + 1 + v2 = v0 + grid_size_y + v3 = v2 + 1 + + indices = np.empty((4 * n_quad + 2 * n_side, 3), dtype=np.int32) + + # Top surface faces (counter-clockwise), two triangles per cell. + top_faces = indices[: 2 * n_quad].reshape(n_quad, 2, 3) + top_faces[:, 0, 0] = v0 + top_faces[:, 0, 1] = v2 + top_faces[:, 0, 2] = v1 + top_faces[:, 1, 0] = v1 + top_faces[:, 1, 1] = v2 + top_faces[:, 1, 2] = v3 + + # Bottom surface faces (clockwise), offset onto the bottom vertex block. + bottom_faces = indices[2 * n_quad : 4 * n_quad].reshape(n_quad, 2, 3) + bottom_faces[:, 0, 0] = v0 + bottom_faces[:, 0, 1] = v1 + bottom_faces[:, 0, 2] = v2 + bottom_faces[:, 1, 0] = v1 + bottom_faces[:, 1, 1] = v3 + bottom_faces[:, 1, 2] = v2 + bottom_faces += num_top_vertices + + # Side walls, one strip per boundary edge, written in place after the surfaces. + i_edge = np.arange(grid_size_x - 1, dtype=np.int32) + j_edge = np.arange(grid_size_y - 1, dtype=np.int32) + side = indices[4 * n_quad :].reshape(n_side, 2, 3) + + def _write_wall(offset: int, t0, t1, flip: bool): + """Emit one quad strip between top edge (t0, t1) and its bottom counterpart.""" + b0 = t0 + num_top_vertices + b1 = t1 + num_top_vertices + wall = side[offset : offset + len(t0)] + if flip: + wall[:, 0, 0], wall[:, 0, 1], wall[:, 0, 2] = t0, t1, b0 + wall[:, 1, 0], wall[:, 1, 1], wall[:, 1, 2] = t1, b1, b0 + else: + wall[:, 0, 0], wall[:, 0, 1], wall[:, 0, 2] = t0, b0, t1 + wall[:, 1, 0], wall[:, 1, 1], wall[:, 1, 2] = t1, b0, b1 + return offset + len(t0) + + offset = 0 # Front edge (j=0) - t0, t1 = i_edge * grid_size_y, (i_edge + 1) * grid_size_y - b0, b1 = num_top_vertices + t0, num_top_vertices + t1 - side_faces_list.append( - np.column_stack([np.column_stack([t0, b0, t1]), np.column_stack([t1, b0, b1])]).reshape(-1, 3) - ) - + offset = _write_wall(offset, i_edge * grid_size_y, (i_edge + 1) * grid_size_y, flip=False) # Back edge (j=grid_size_y-1) - t0 = i_edge * grid_size_y + (grid_size_y - 1) - t1 = (i_edge + 1) * grid_size_y + (grid_size_y - 1) - b0, b1 = num_top_vertices + t0, num_top_vertices + t1 - side_faces_list.append( - np.column_stack([np.column_stack([t0, t1, b0]), np.column_stack([t1, b1, b0])]).reshape(-1, 3) + offset = _write_wall( + offset, + i_edge * grid_size_y + (grid_size_y - 1), + (i_edge + 1) * grid_size_y + (grid_size_y - 1), + flip=True, ) - # Left edge (i=0) - t0, t1 = j_edge, j_edge + 1 - b0, b1 = num_top_vertices + t0, num_top_vertices + t1 - side_faces_list.append( - np.column_stack([np.column_stack([t0, t1, b0]), np.column_stack([t1, b1, b0])]).reshape(-1, 3) - ) - + offset = _write_wall(offset, j_edge, j_edge + 1, flip=True) # Right edge (i=grid_size_x-1) - t0 = (grid_size_x - 1) * grid_size_y + j_edge - t1 = (grid_size_x - 1) * grid_size_y + (j_edge + 1) - b0, b1 = num_top_vertices + t0, num_top_vertices + t1 - side_faces_list.append( - np.column_stack([np.column_stack([t0, b0, t1]), np.column_stack([t1, b0, b1])]).reshape(-1, 3) + _write_wall( + offset, + (grid_size_x - 1) * grid_size_y + j_edge, + (grid_size_x - 1) * grid_size_y + (j_edge + 1), + flip=False, ) - # Combine all faces - all_faces = np.vstack([top_faces, bottom_faces, *side_faces_list]) - indices = all_faces.astype(np.int32).flatten() + indices = indices.reshape(-1) return vertices, indices diff --git a/newton/_src/geometry/types.py b/newton/_src/geometry/types.py index a559f15033..3ce5362af1 100644 --- a/newton/_src/geometry/types.py +++ b/newton/_src/geometry/types.py @@ -2211,6 +2211,52 @@ def __init__( self.mass = 0.0 self.com = wp.vec3() + @staticmethod + def create_from_mesh( + mesh: "wp.Mesh", + resolution: float, + *, + max_cells_per_axis: int = 4096, + ) -> tuple["Heightfield", wp.transform]: + """Create a heightfield by rasterizing a triangle mesh. + + Rays are cast straight down onto the mesh on a regular grid to sample its + elevation (see :func:`~newton.utils.rasterize_mesh_to_heightfield`). This + method supports terrain that is single-valued in Z, including sloped planes. + + Args: + mesh: Triangle mesh to rasterize, with vertex coordinates [m] in the + frame where it should be placed. + resolution: Horizontal grid spacing [m]. Smaller values preserve more + detail at the cost of a larger grid. + max_cells_per_axis: Upper bound on grid rows/columns. If the mesh extent + would exceed this, the effective resolution is coarsened to fit. + + Returns: + A tuple ``(heightfield, xform)`` where ``heightfield`` is the sampled + :class:`Heightfield` and ``xform`` has a translation [m] that places its + (origin-centered) grid at the mesh's XY center. Pass both to + :meth:`~newton.ModelBuilder.add_shape_heightfield`. + """ + from ..utils.heightfield import rasterize_mesh_to_heightfield # noqa: PLC0415 + + heights, (x_min, y_min, x_max, y_max) = rasterize_mesh_to_heightfield( + mesh, resolution, max_cells_per_axis=max_cells_per_axis + ) + nrow, ncol = heights.shape + heightfield = Heightfield( + data=heights, + nrow=nrow, + ncol=ncol, + hx=0.5 * (x_max - x_min), + hy=0.5 * (y_max - y_min), + ) + xform = wp.transform( + wp.vec3(0.5 * (x_min + x_max), 0.5 * (y_min + y_max), 0.0), + wp.quat_identity(), + ) + return heightfield, xform + @property def data(self): """Get the normalized [0, 1] elevation data as a 2D numpy array.""" diff --git a/newton/_src/sensors/sensor_contact.py b/newton/_src/sensors/sensor_contact.py index dfe235c61a..8c51f8a38f 100644 --- a/newton/_src/sensors/sensor_contact.py +++ b/newton/_src/sensors/sensor_contact.py @@ -3,6 +3,7 @@ from __future__ import annotations +import re import warnings from typing import Any, Literal @@ -452,10 +453,10 @@ def __init__( self, model: Model, *, - sensing_bodies: str | list[str] | list[int] | None = None, - sensing_shapes: str | list[str] | list[int] | None = None, - counterpart_bodies: str | list[str] | list[int] | None = None, - counterpart_shapes: str | list[str] | list[int] | None = None, + sensing_bodies: str | list[str] | re.Pattern[str] | list[int] | None = None, + sensing_shapes: str | list[str] | re.Pattern[str] | list[int] | None = None, + counterpart_bodies: str | list[str] | re.Pattern[str] | list[int] | None = None, + counterpart_shapes: str | list[str] | re.Pattern[str] | list[int] | None = None, measure_total: bool = True, verbose: bool | None = None, request_contact_attributes: bool = True, @@ -470,14 +471,14 @@ def __init__( Args: model: The simulation model providing shape/body definitions and world layout. - sensing_bodies: List of body indices, single pattern to match against body labels, or list of patterns where - any one matches. - sensing_shapes: List of shape indices, single pattern to match against shape labels, or list of patterns - where any one matches. - counterpart_bodies: List of body indices, single pattern to match - against body labels, or list of patterns where any one matches. - counterpart_shapes: List of shape indices, single pattern to match - against shape labels, or list of patterns where any one matches. + sensing_bodies: Glob pattern, list of glob patterns, compiled regular-expression pattern to match against + body labels, or list of body indices. Regular expressions use full matching. + sensing_shapes: Glob pattern, list of glob patterns, compiled regular-expression pattern to match against + shape labels, or list of shape indices. Regular expressions use full matching. + counterpart_bodies: Glob pattern, list of glob patterns, compiled regular-expression pattern to match + against body labels, or list of body indices. Regular expressions use full matching. + counterpart_shapes: Glob pattern, list of glob patterns, compiled regular-expression pattern to match + against shape labels, or list of shape indices. Regular expressions use full matching. measure_total: If True (default), :attr:`total_force` and :attr:`total_force_friction` are allocated. If False, both are None. verbose: If True, print details. If False, suppress details. If None, print details when diff --git a/newton/_src/sensors/sensor_frame_transform.py b/newton/_src/sensors/sensor_frame_transform.py index 6f9d5e7b60..b11c81e1e3 100644 --- a/newton/_src/sensors/sensor_frame_transform.py +++ b/newton/_src/sensors/sensor_frame_transform.py @@ -3,6 +3,8 @@ """Frame Transform Sensor - measures transforms relative to sites.""" +import re + import warp as wp from ..geometry import ShapeFlags @@ -120,8 +122,8 @@ class SensorFrameTransform: def __init__( self, model: Model, - shapes: str | list[str] | list[int], - reference_sites: str | list[str] | list[int], + shapes: str | list[str] | re.Pattern[str] | list[int], + reference_sites: str | list[str] | re.Pattern[str] | list[int], *, verbose: bool | None = None, ): @@ -129,10 +131,11 @@ def __init__( Args: model: The model to measure. - shapes: List of shape indices, single pattern to match against shape - labels, or list of patterns where any one matches. - reference_sites: List of site indices, single pattern to match against - site labels, or list of patterns where any one matches. Must expand + shapes: Glob pattern, list of glob patterns, compiled regular-expression + pattern to match against shape labels, or list of shape indices. Regular + expressions use full matching. + reference_sites: Glob pattern, list of glob patterns, compiled regular-expression + pattern to match against site labels, or list of site indices. Must expand to one site or the same number as ``shapes``. verbose: If True, print details. If False, suppress details. If None, print details when ``wp.config.log_level`` is configured for debug logging. diff --git a/newton/_src/sensors/sensor_imu.py b/newton/_src/sensors/sensor_imu.py index 63f807cd75..6a9e70e051 100644 --- a/newton/_src/sensors/sensor_imu.py +++ b/newton/_src/sensors/sensor_imu.py @@ -3,6 +3,8 @@ """IMU Sensor - measures accelerations and angular velocities at sensor sites.""" +import re + import warp as wp from ..geometry.flags import ShapeFlags @@ -114,7 +116,7 @@ class SensorIMU: def __init__( self, model: Model, - sites: str | list[str] | list[int], + sites: str | list[str] | re.Pattern[str] | list[int], *, verbose: bool | None = None, request_state_attributes: bool = True, @@ -126,8 +128,9 @@ def __init__( Args: model: The model to use. - sites: List of site indices, single pattern to match against site - labels, or list of patterns where any one matches. + sites: Glob pattern, list of glob patterns, compiled regular-expression + pattern to match against site labels, or list of site indices. Regular + expressions use full matching. verbose: If True, print details. If False, suppress details. If None, print details when ``wp.config.log_level`` is configured for debug logging. request_state_attributes: If True (default), transparently request the extended state attribute ``body_qdd`` from the model. diff --git a/newton/_src/sensors/sensor_tiled_camera.py b/newton/_src/sensors/sensor_tiled_camera.py index da60d444cc..397be6a0bc 100644 --- a/newton/_src/sensors/sensor_tiled_camera.py +++ b/newton/_src/sensors/sensor_tiled_camera.py @@ -3,6 +3,7 @@ from __future__ import annotations +import os import warnings from typing import Any @@ -19,6 +20,8 @@ Utils, ) +PROFILE_ENABLED = os.environ.get("NEWTON_PROFILE", "0") != "0" + _RENDER_CONFIG_DEPRECATION_MSG = ( "SensorTiledCamera.render_config is deprecated as of Newton 1.4; " "use SensorTiledCamera.default_render_config instead. " @@ -224,24 +227,27 @@ def update( for the render megakernel. """ - self.sync_transforms(state) - - self.__render_context.render( - self.model, - state, - camera_transforms=camera_transforms, - camera_rays=camera_rays, - color_image=color_image, - hdr_color_image=hdr_color_image, - depth_image=depth_image, - forward_depth_image=forward_depth_image, - shape_index_image=shape_index_image, - normal_image=normal_image, - albedo_image=albedo_image, - clear_data=clear_data if clear_data is not None else self.default_clear_data, - config=render_config if render_config is not None else self.default_render_config, - kernel_block_dim=kernel_block_dim, - ) + with wp.ScopedTimer( + "Newton::SensorTiledCamera::update", active=PROFILE_ENABLED, use_nvtx=True, synchronize=True + ): + self.sync_transforms(state) + + self.__render_context.render( + self.model, + state, + camera_transforms=camera_transforms, + camera_rays=camera_rays, + color_image=color_image, + hdr_color_image=hdr_color_image, + depth_image=depth_image, + forward_depth_image=forward_depth_image, + shape_index_image=shape_index_image, + normal_image=normal_image, + albedo_image=albedo_image, + clear_data=clear_data if clear_data is not None else self.default_clear_data, + config=render_config if render_config is not None else self.default_render_config, + kernel_block_dim=kernel_block_dim, + ) @property def render_config(self) -> RenderConfig: diff --git a/newton/_src/sensors/warp_raytrace/render.py b/newton/_src/sensors/warp_raytrace/render.py index 9186e577d5..b2694d1969 100644 --- a/newton/_src/sensors/warp_raytrace/render.py +++ b/newton/_src/sensors/warp_raytrace/render.py @@ -64,9 +64,9 @@ def write_clear_outputs( out_hdr_color: wp.array[wp.vec3f], ): if wp.static(state.render_color): - out_color[out_index] = wp.uint32(wp.static(clear_data.clear_color)) + out_color[out_index] = wp.static(wp.uint32(clear_data.clear_color)) if wp.static(state.render_albedo): - out_albedo[out_index] = wp.uint32(wp.static(clear_data.clear_albedo)) + out_albedo[out_index] = wp.static(wp.uint32(clear_data.clear_albedo)) if wp.static(state.render_hdr_color): out_hdr_color[out_index] = wp.vec3f(0.0) if wp.static(state.render_depth): @@ -80,7 +80,7 @@ def write_clear_outputs( wp.static(clear_data.clear_normal[2]), ) if wp.static(state.render_shape_index): - out_shape_index[out_index] = wp.uint32(wp.static(clear_data.clear_shape_index)) + out_shape_index[out_index] = wp.static(wp.uint32(clear_data.clear_shape_index)) @wp.kernel(enable_backward=False, module="unique", module_options={"fast_math": config.enable_fast_math}) def render_megakernel( diff --git a/newton/_src/sensors/warp_raytrace/utils.py b/newton/_src/sensors/warp_raytrace/utils.py index a0e3360892..f6d8b84f4d 100644 --- a/newton/_src/sensors/warp_raytrace/utils.py +++ b/newton/_src/sensors/warp_raytrace/utils.py @@ -18,6 +18,10 @@ if TYPE_CHECKING: from .render_context import RenderContext +# Knuth multiplicative hash constant (2^32 / golden ratio). +# Typed uint32 so kernel codegen doesn't overflow an int32 constant. +HASH_MULTIPLIER = wp.uint32(2654435761) + def _resolve_fisheye_image_size( axis: str, @@ -225,7 +229,7 @@ def unpack_shape_index_hash_to_rgba_kernel( # Knuth multiplicative hash, masked to 24 bits. ``idx + 1`` keeps shape 0 # away from the all-zero hash that collides with the miss color; the # miss sentinel ``0xFFFFFFFF`` wraps back to 0 and intentionally renders black. - h = ((idx + wp.uint32(1)) * wp.uint32(2654435761)) & wp.uint32(0xFFFFFF) + h = ((idx + wp.uint32(1)) * HASH_MULTIPLIER) & wp.uint32(0xFFFFFF) out[n, y, x, 0] = wp.uint8((h >> wp.uint32(16)) & wp.uint32(0xFF)) out[n, y, x, 1] = wp.uint8((h >> wp.uint32(8)) & wp.uint32(0xFF)) out[n, y, x, 2] = wp.uint8(h & wp.uint32(0xFF)) diff --git a/newton/_src/sim/articulation.py b/newton/_src/sim/articulation.py index ae2cdec7cd..82a9037d1a 100644 --- a/newton/_src/sim/articulation.py +++ b/newton/_src/sim/articulation.py @@ -631,8 +631,7 @@ def reconstruct_angular_q_qd(q_pc: wp.quat, w_err: wp.vec3, X_wp: wp.transform, qd: The joint velocity coordinate. """ axis_p = wp.transform_vector(X_wp, axis) - twist = wp.quat_twist(axis, q_pc) - q = wp.acos(twist[3]) * 2.0 * wp.sign(wp.dot(axis, wp.vec3(twist[0], twist[1], twist[2]))) + q = wp.quat_twist_angle_signed(axis, q_pc) qd = wp.dot(w_err, axis_p) return q, qd diff --git a/newton/_src/sim/builder.py b/newton/_src/sim/builder.py index bb62a306d7..eb752b626e 100644 --- a/newton/_src/sim/builder.py +++ b/newton/_src/sim/builder.py @@ -11,6 +11,7 @@ import inspect import math import warnings +import weakref from collections import Counter, deque from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass, replace @@ -89,6 +90,13 @@ _IDENTITY_TRANSFORM = np.asarray(wp.transformf((0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0)), dtype=np.float32) _IDENTITY_ROTATION = np.asarray(wp.quatf(0.0, 0.0, 0.0, 1.0), dtype=np.float32) +_MERGE_VALIDATION_CACHE: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() +"""Memoizes :meth:`ModelBuilder._validate_builder_merge` as ``dest -> {source: schema epochs}``. + +Kept out of the builders themselves: this is a memoization table, not model state, and it must +not show up in builder-state comparisons or survive past either builder's lifetime. +""" + @dataclass(frozen=True) class _ShapeCollisionFilterBlock: @@ -424,6 +432,9 @@ class BvhConfig: shape_constructor: str | None = None """Warp model shape BVH constructor backend. If ``None``, Warp's default is used.""" + shape_flags: ShapeFlags = ShapeFlags.VISIBLE + """Mask of :class:`~newton.ShapeFlags`; a shape is included in the model shape BVH if any of its flags are set in the mask.""" + @dataclass class MeshApproximationConfig: """Default settings for mesh approximation. @@ -492,7 +503,7 @@ class ShapeConfig: """Indicates whether the shape is visible in the simulation. Defaults to True.""" is_site: bool = False """Indicates whether the shape is a site (non-colliding reference point). Directly setting this to True will NOT enforce site invariants. Use `mark_as_site()` or set via the `flags` property to ensure invariants. Defaults to False.""" - sdf_narrow_band_range: tuple[float, float] = (-0.1, 0.1) + sdf_narrow_band_range: tuple[float, float] | list[float] = (-0.1, 0.1) """The narrow band distance range (inner, outer) for primitive SDF computation.""" sdf_target_voxel_size: float | None = None """Target voxel size for sparse SDF grid. @@ -601,13 +612,40 @@ def validate(self, shape_type: int | None = None) -> None: raise ValueError( f"Unknown sdf_texture_format {self.sdf_texture_format!r}. Expected one of {list(_valid_tex_fmts)}." ) - if self.sdf_max_resolution is not None and self.sdf_target_voxel_size is not None: - raise ValueError("Set only one of sdf_max_resolution or sdf_target_voxel_size, not both.") - if self.sdf_max_resolution is not None and self.sdf_max_resolution % 8 != 0: + if not math.isfinite(self.density) or self.density < 0.0: + raise ValueError(f"density must be finite and >= 0 (got {self.density}).") + + if self.sdf_target_voxel_size is not None and ( + not math.isfinite(self.sdf_target_voxel_size) or self.sdf_target_voxel_size <= 0.0 + ): + raise ValueError(f"sdf_target_voxel_size must be finite and > 0 (got {self.sdf_target_voxel_size}).") + + if self.sdf_padding is not None and (not math.isfinite(self.sdf_padding) or self.sdf_padding < 0.0): + raise ValueError(f"sdf_padding must be finite and >= 0 (got {self.sdf_padding}).") + + if not isinstance(self.sdf_narrow_band_range, (tuple, list)) or len(self.sdf_narrow_band_range) != 2: + raise ValueError( + "sdf_narrow_band_range must contain two distances (inner, outer) with inner < 0 < outer." + ) + inner, outer = self.sdf_narrow_band_range + if not math.isfinite(inner) or not math.isfinite(outer) or not inner < 0.0 < outer: raise ValueError( - f"sdf_max_resolution must be divisible by 8 (got {self.sdf_max_resolution}). " - "This is required because SDF volumes are allocated in 8x8x8 tiles." + f"sdf_narrow_band_range must contain finite values satisfying inner < 0 < outer " + f"(got {self.sdf_narrow_band_range})." ) + + if self.sdf_max_resolution is not None and self.sdf_target_voxel_size is not None: + raise ValueError("Set only one of sdf_max_resolution or sdf_target_voxel_size, not both.") + if self.sdf_max_resolution is not None: + if self.sdf_max_resolution <= 0: + raise ValueError(f"sdf_max_resolution must be > 0 (got {self.sdf_max_resolution}).") + if self.sdf_max_resolution >= (1 << 16): + raise ValueError(f"sdf_max_resolution must be less than {1 << 16}.") + if self.sdf_max_resolution % 8 != 0: + raise ValueError( + f"sdf_max_resolution must be divisible by 8 (got {self.sdf_max_resolution}). " + "This is required because SDF volumes are allocated in 8x8x8 tiles." + ) hydroelastic_supported = shape_type not in (GeoType.PLANE, GeoType.HFIELD) hydroelastic_requires_configured_sdf = shape_type in ( GeoType.SPHERE, @@ -1252,6 +1290,10 @@ def __init__( self._shape_collision_filter_pairs: _BuilderShapeCollisionFilterPairs | list[tuple[int, int]] = ( _BuilderShapeCollisionFilterPairs() ) + self._merge_filter_template: tuple[list[tuple[int, int]], int, tuple[tuple[int, int], ...]] | None = None + """Cache backing :meth:`_materialized_filter_template`, as ``(source, length, template)``.""" + self._custom_schema_epoch: int = 0 + """Bumped whenever the custom attribute/frequency registry changes; keys the merge-validation cache.""" self._requested_contact_attributes: set[str] = set() """Optional contact attributes requested via :meth:`request_contact_attributes`.""" @@ -1609,6 +1651,25 @@ def shape_collision_filter_pairs(self) -> list[tuple[int, int]]: def shape_collision_filter_pairs(self, pairs: list[tuple[int, int]]) -> None: self._shape_collision_filter_pairs = pairs + def _materialized_filter_template(self) -> tuple[tuple[int, int], ...]: + """This builder's filter pairs as one tuple, stable across repeated merges. + + Merging a source builder world-by-world (one :meth:`add_builder` per world) must + hand out the *same* tuple object every time: the collision-filter and + contact-pair template caches in :meth:`finalize` are keyed by object identity, so + a freshly built tuple per world silently disables them and makes finalization + scale with world count instead of with the number of distinct sources. + """ + pairs = self._shape_collision_filter_pairs + if isinstance(pairs, _BuilderShapeCollisionFilterPairs): + return pairs.template_pairs() + cached = self._merge_filter_template + if cached is not None and cached[0] is pairs and cached[1] == len(pairs): + return cached[2] + template = tuple(pairs) + self._merge_filter_template = (pairs, len(pairs), template) + return template + def add_shape_collision_filter_pair(self, shape_a: int, shape_b: int) -> None: """Add a collision filter pair in canonical order. @@ -1707,6 +1768,7 @@ def add_custom_attribute(self, attribute: CustomAttribute) -> None: ) self.custom_attributes[key] = attribute + self._custom_schema_epoch += 1 def _add_custom_attribute_model_finalizer( self, @@ -1760,6 +1822,7 @@ def add_custom_frequency(self, frequency: CustomFrequency) -> None: return self.custom_frequencies[freq_key] = freq_obj + self._custom_schema_epoch += 1 if freq_key not in self._custom_frequency_counts: self._custom_frequency_counts[freq_key] = 0 @@ -2606,6 +2669,8 @@ def replicate( builder: ModelBuilder, world_count: int, spacing: tuple[float, float, float] = (0.0, 0.0, 0.0), + *, + xforms: Sequence[Transform] | None = None, ): """ Replicates the given builder multiple times, offsetting each copy according to the supplied spacing. @@ -2634,9 +2699,12 @@ def replicate( Args: builder: The builder to replicate. All entities from this builder will be copied. world_count: The number of worlds to create. - spacing: The spacing between each copy along each axis. + spacing: The spacing between each copy along each axis. Ignored when + ``xforms`` is provided. For example, (5.0, 5.0, 0.0) arranges copies in a 2D grid in the XY plane. Defaults to (0.0, 0.0, 0.0). + xforms: Optional sequence of transforms, one per replicated world. + When provided, its length must equal ``world_count``. """ if world_count <= 0: return @@ -2645,10 +2713,14 @@ def replicate( f"Cannot begin a new world: already in world context (current_world={self.current_world}). " "Call end_world() first to close the current world context." ) - offsets = compute_world_offsets(world_count, spacing, self.up_axis) + if xforms is None: + offsets = compute_world_offsets(world_count, spacing, self.up_axis) + xforms = [wp.transform(offset, wp.quat_identity()) for offset in offsets] + elif len(xforms) != world_count: + raise ValueError(f"xforms must contain {world_count} entries, got {len(xforms)}") + base_world = self.world_count worlds = list(range(base_world, base_world + world_count)) - xforms = [wp.transform(offset, wp.quat_identity()) for offset in offsets] self._merge_builder_copies(builder, worlds, xforms, [None] * world_count) self.world_gravity.extend(builder._gravity_as_vector() for _ in range(world_count)) @@ -2828,11 +2900,7 @@ def extend_referenced(dst: list, values: Sequence[Any], kind: str) -> None: source_filter_pairs = builder._shape_collision_filter_pairs if source_filter_pairs: - template_pairs = ( - source_filter_pairs.template_pairs() - if isinstance(source_filter_pairs, _BuilderShapeCollisionFilterPairs) - else tuple(source_filter_pairs) - ) + template_pairs = builder._materialized_filter_template() for world, shape_start in zip(worlds.tolist(), shape_starts.tolist(), strict=True): if isinstance(self._shape_collision_filter_pairs, _BuilderShapeCollisionFilterPairs): self._shape_collision_filter_pairs.extend_offset( @@ -3011,6 +3079,18 @@ def _custom_attribute_defaults_match(existing: Any, incoming: Any) -> bool: return bool(matches) def _validate_builder_merge(self, builder: ModelBuilder, entity_kinds: set[str]) -> None: + # Replication merges the same handful of source builders once per world, and this + # check is pure schema validation: it compares custom-attribute specs and defaults, + # which cannot change unless one of the two registries changes. Both are versioned, + # so a repeat merge of an unchanged pair is skipped. Without this the element-wise + # equality on Warp-typed defaults runs once per attribute per world. + # ``valid_references`` only ever grows as merges accumulate frequency counts, so a + # pair that validated before still validates. + cache_key = (self._custom_schema_epoch, builder._custom_schema_epoch, frozenset(entity_kinds)) + validated = _MERGE_VALIDATION_CACHE.setdefault(self, weakref.WeakKeyDictionary()) + if validated.get(builder) == cache_key: + return + valid_references = entity_kinds | set(self._custom_frequency_counts) | set(builder._custom_frequency_counts) for freq_key, frequency in builder.custom_frequencies.items(): @@ -3059,6 +3139,8 @@ def _validate_builder_merge(self, builder: ModelBuilder, entity_kinds: set[str]) f"({existing!r} != {finalizer!r})." ) + validated[builder] = cache_key + def add_articulation( self, joints: list[int], label: str | None = None, custom_attributes: dict[str, Any] | None = None ): @@ -3381,6 +3463,7 @@ def add_usd( skip_mesh_approximation: bool = False, load_sites: bool = True, load_visual_shapes: bool = True, + load_static_visual_shapes: bool = True, hide_collision_shapes: bool = False, force_show_colliders: bool = False, parse_mujoco_options: bool = True, @@ -3487,6 +3570,9 @@ def add_usd( skip_mesh_approximation: If True, mesh approximation is skipped. Otherwise, meshes are approximated according to the ``physics:approximation`` attribute defined on the UsdPhysicsMeshCollisionAPI (if it is defined), using the settings from :attr:`~newton.ModelBuilder.default_mesh_approximation_cfg`. Default is False. load_sites: If True, sites (prims with ``NewtonSiteAPI`` or ``MjcSiteAPI``) are loaded as non-colliding reference points. If False, sites are ignored. Default is True. load_visual_shapes: If True, non-physics visual geometry is loaded. If False, visual-only shapes are ignored (sites are still controlled by ``load_sites``). Default is True. + load_static_visual_shapes: If True, supported visual-only geometry outside + rigid-body hierarchies is loaded as static shapes when + ``load_visual_shapes`` is also True. Default is True. hide_collision_shapes: If True, collision shapes on bodies that already have visual-only geometry are hidden unconditionally, regardless of whether the collider has authored PBR material data. Default is False. @@ -3625,6 +3711,7 @@ def add_usd( skip_mesh_approximation=skip_mesh_approximation, load_sites=load_sites, load_visual_shapes=load_visual_shapes, + load_static_visual_shapes=load_static_visual_shapes, hide_collision_shapes=hide_collision_shapes, force_show_colliders=force_show_colliders, parse_mujoco_options=parse_mujoco_options, @@ -4008,6 +4095,7 @@ def get_offset(entity_or_key: str | None) -> int: freq_key = attr.frequency mapped_values = [] if isinstance(freq_key, str) else {} self.custom_attributes[full_key] = replace(attr, values=mapped_values) + self._custom_schema_epoch += 1 continue freq_key = attr.frequency @@ -4099,6 +4187,7 @@ def transform_enum_value( else: mapped_values = {index_offset + idx: value for idx, value in attr.values.items()} self.custom_attributes[full_key] = replace(attr, values=mapped_values) + self._custom_schema_epoch += 1 continue if not self._custom_attribute_defaults_match(merged.default, attr.default): @@ -4137,6 +4226,7 @@ def transform_enum_value( for freq_key, freq_obj in builder.custom_frequencies.items(): if freq_key not in self.custom_frequencies: self.custom_frequencies[freq_key] = freq_obj + self._custom_schema_epoch += 1 for freq_key, builder_count in builder._custom_frequency_counts.items(): offset = custom_frequency_offsets.get(freq_key, 0) @@ -4817,6 +4907,7 @@ def add_joint_ball( child_xform: Transform | None = None, armature: float | None = None, friction: float | None = None, + damping: float | None = None, label: str | None = None, collision_filter_parent: bool | None = None, enabled: bool = True, @@ -4832,6 +4923,7 @@ def add_joint_ball( child_xform: The transform from the child body frame to the joint child anchor frame. armature: Artificial inertia added around the joint axes. If None, the default value from ``ModelBuilder.default_joint_cfg.armature`` is used. friction: Friction coefficient for the joint axes. If None, the default value from ``ModelBuilder.default_joint_cfg.friction`` is used. + damping: Passive angular velocity damping [N·s/m or N·m·s/rad, depending on joint type] always active on all three BALL joint angular DOFs. If None, the default value from ``ModelBuilder.default_joint_cfg.damping`` is used. label: The label of the joint. collision_filter_parent: Whether to filter collisions between shapes of the parent and child bodies. Defaults to ``False`` for joints to world, ``True`` otherwise. enabled: Whether the joint is enabled. @@ -4849,23 +4941,28 @@ def add_joint_ball( armature = self.default_joint_cfg.armature if friction is None: friction = self.default_joint_cfg.friction + if damping is None: + damping = self.default_joint_cfg.damping x = ModelBuilder.JointDofConfig( axis=Axis.X, armature=armature, friction=friction, + damping=damping, actuator_mode=actuator_mode, ) y = ModelBuilder.JointDofConfig( axis=Axis.Y, armature=armature, friction=friction, + damping=damping, actuator_mode=actuator_mode, ) z = ModelBuilder.JointDofConfig( axis=Axis.Z, armature=armature, friction=friction, + damping=damping, actuator_mode=actuator_mode, ) @@ -5120,40 +5217,68 @@ def add_joint_cable( child_xform: Transform | None = None, stretch_stiffness: float | None = None, stretch_damping: float | None = None, + shear_stiffness: float | None = None, + shear_damping: float | None = None, bend_stiffness: float | None = None, bend_damping: float | None = None, + twist_stiffness: float | None = None, + twist_damping: float | None = None, label: str | None = None, collision_filter_parent: bool | None = None, enabled: bool = True, custom_attributes: dict[str, Any] | None = None, **kwargs, ) -> int: - """Adds a cable joint to the model. It has two degrees of freedom: one linear (stretch) - that constrains the distance between the attachment points, and one angular (bend/twist) - that penalizes the relative rotation of the attachment frames. + """Adds a cable joint to the model. + + Cable joints have split linear stretch/shear DoFs plus separate angular + bend and twist DoFs. When both ``shear_stiffness`` and + ``shear_damping`` are omitted, shear uses the stretch stiffness / + damping, reproducing the isotropic linear energy while using the + split layout. When both ``twist_stiffness`` and ``twist_damping`` are + omitted, twist uses the bend stiffness / damping, reproducing the + isotropic angular energy while using the split layout. .. note:: - Cable joints are represented in the joint data model, but their two entries - are VBD stretch and bend/twist constraint slots rather than - ``joint_q`` coordinates. Cable body transforms are integrated directly by - :class:`newton.solvers.SolverVBD`; they are not reconstructed by - :func:`newton.eval_fk`. + Cable joints are supported by :class:`newton.solvers.SolverVBD`, which uses an + AVBD backend for rigid bodies. Split cables are represented in the + joint data model as VBD stretch, shear, bend, and twist constraint + slots rather than ``joint_q`` coordinates. Cable body transforms are + integrated directly by :class:`newton.solvers.SolverVBD`; they are + not reconstructed by :func:`newton.eval_fk`. + + Split cables use each anchor frame's local ``+Z`` as the material + tangent axis for separating axial stretch from shear and twist from + bend. For a body-to-body cable span, the parent anchor ``+Z`` should + point from the parent attachment toward the child attachment. + :meth:`add_rod` and :meth:`add_rod_graph` satisfy the tangent + convention automatically. Args: parent: The index of the parent body. child: The index of the child body. parent_xform: The transform from the parent body frame to the joint parent anchor frame; its - translation is the attachment point. + translation is the attachment point and its local ``+Z`` axis is the parent-side material + tangent. child_xform: The transform from the child body frame to the joint child anchor frame; its - translation is the attachment point. + translation is the attachment point and its local ``+Z`` axis is the child-side material + tangent. stretch_stiffness: Cable stretch stiffness (stored as ``target_ke``) [N/m]. If None, defaults to 1.0e5. stretch_damping: Cable stretch damping [N·s/m] (stored as ``target_kd``). If None, defaults to 0.0. - bend_stiffness: Cable bend/twist stiffness (stored as ``target_ke``) [N*m] (torque per radian). If None, - defaults to 0.0. - bend_damping: Cable bend/twist damping [N·m·s/rad] (stored as ``target_kd``). If None, + shear_stiffness: Optional transverse shear stiffness [N/m]. If None, + defaults to ``stretch_stiffness``. + shear_damping: Optional transverse shear damping [N·s/m]. If None, defaults to + ``stretch_damping`` only when both ``shear_stiffness`` and ``shear_damping`` are None. Otherwise defaults to 0.0. + bend_stiffness: Cable bend stiffness (stored as ``target_ke``) [N*m] + (torque per radian). If None, defaults to 0.0. + bend_damping: Cable bend damping [N·m·s/rad] (stored as ``target_kd``). If None, defaults to 0.0. + twist_stiffness: Optional twist stiffness [N*m] (torque per radian). If None, + defaults to ``bend_stiffness``. + twist_damping: Optional twist damping [N·m·s/rad]. If None, defaults to ``bend_damping`` only when + both ``twist_stiffness`` and ``twist_damping`` are None. Otherwise defaults to 0.0. label: The label of the joint. collision_filter_parent: Whether to filter collisions between shapes of the parent and child bodies. Defaults to ``False`` for joints to world, ``True`` otherwise. enabled: Whether the joint is enabled. @@ -5164,15 +5289,35 @@ def add_joint_cable( The index of the added joint. """ - # Linear DOF (stretch) - se_ke = 1.0e5 if stretch_stiffness is None else stretch_stiffness - se_kd = 0.0 if stretch_damping is None else stretch_damping - ax_lin = ModelBuilder.JointDofConfig(target_ke=se_ke, target_kd=se_kd) + # Linear DOFs (stretch and shear). Default shear to stretch so omitted + # shear reproduces the isotropic linear anchor energy in the split layout. + stretch_ke = 1.0e5 if stretch_stiffness is None else stretch_stiffness + stretch_kd = 0.0 if stretch_damping is None else stretch_damping + stretch_axis = ModelBuilder.JointDofConfig(target_ke=stretch_ke, target_kd=stretch_kd) + if shear_stiffness is None and shear_damping is None: + shear_ke = stretch_ke + shear_kd = stretch_kd + else: + shear_ke = stretch_ke if shear_stiffness is None else shear_stiffness + shear_kd = 0.0 if shear_damping is None else shear_damping + shear_axis = ModelBuilder.JointDofConfig(target_ke=shear_ke, target_kd=shear_kd) - # Angular DOF (bend/twist) + # Angular DOFs (bend and twist). Default twist to bend so omitted twist + # reproduces the isotropic angular energy in the split layout. bend_ke = 0.0 if bend_stiffness is None else bend_stiffness bend_kd = 0.0 if bend_damping is None else bend_damping - ax_ang = ModelBuilder.JointDofConfig(target_ke=bend_ke, target_kd=bend_kd) + bend_axis = ModelBuilder.JointDofConfig(target_ke=bend_ke, target_kd=bend_kd) + if twist_stiffness is None and twist_damping is None: + twist_ke = bend_ke + twist_kd = bend_kd + else: + twist_ke = bend_ke if twist_stiffness is None else twist_stiffness + twist_kd = 0.0 if twist_damping is None else twist_damping + if stretch_ke < 0.0 or shear_ke < 0.0 or bend_ke < 0.0 or twist_ke < 0.0: + raise ValueError( + "add_joint_cable: stretch_stiffness, shear_stiffness, bend_stiffness, and twist_stiffness must be >= 0" + ) + twist_axis = ModelBuilder.JointDofConfig(target_ke=twist_ke, target_kd=twist_kd) return self.add_joint( JointType.CABLE, @@ -5180,8 +5325,8 @@ def add_joint_cable( child, parent_xform=parent_xform, child_xform=child_xform, - linear_axes=[ax_lin], - angular_axes=[ax_ang], + linear_axes=[stretch_axis, shear_axis], + angular_axes=[bend_axis, twist_axis], label=label, collision_filter_parent=collision_filter_parent, enabled=enabled, @@ -7506,8 +7651,12 @@ def add_rod( cfg: ShapeConfig | None = None, stretch_stiffness: float | None = None, stretch_damping: float | None = None, + shear_stiffness: float | None = None, + shear_damping: float | None = None, bend_stiffness: float | None = None, bend_damping: float | None = None, + twist_stiffness: float | None = None, + twist_damping: float | None = None, closed: bool = False, label: str | None = None, wrap_in_articulation: bool = True, @@ -7518,8 +7667,8 @@ def add_rod( Constructs a chain of capsule bodies from the given centerline points and orientations. Each segment is a capsule aligned by the corresponding quaternion, and adjacent capsules - are connected by cable joints providing one linear (stretch) and one angular (bend/twist) - degree of freedom. + are connected by cable joints providing split linear stretch/shear and split angular + bend/twist degrees of freedom. Args: positions: Centerline node positions (segment endpoints) in world space. These are the @@ -7535,10 +7684,18 @@ def add_rod( If None, defaults to 1.0e5. stretch_damping: Stretch damping [N·s/m] for the cable joints (applied per-joint; not length-normalized). If None, defaults to 0.0. - bend_stiffness: Per-joint cable bend/twist stiffness, stored directly as ``target_ke`` + shear_stiffness: Optional per-joint transverse shear stiffness [N/m]. If None, defaults to + ``stretch_stiffness``. + shear_damping: Optional per-joint transverse shear damping [N·s/m]. If None, defaults to + ``stretch_damping`` only when both ``shear_stiffness`` and ``shear_damping`` are None. Otherwise defaults to 0.0. + bend_stiffness: Per-joint cable bend stiffness, stored directly as ``target_ke`` [N*m] (torque per radian). If None, defaults to 0.0. - bend_damping: Bend/twist damping [N·m·s/rad] for the cable joints (applied per-joint; not length-normalized). If None, + bend_damping: Bend damping [N·m·s/rad] for the cable joints (applied per-joint; not length-normalized). If None, defaults to 0.0. + twist_stiffness: Optional per-joint cable twist stiffness [N*m]. If None, defaults to + ``bend_stiffness``. + twist_damping: Optional per-joint cable twist damping [N·m·s/rad]. If None, defaults to ``bend_damping`` + only when both ``twist_stiffness`` and ``twist_damping`` are None. Otherwise defaults to 0.0. closed: If True, connects the last segment back to the first to form a closed loop. If False, creates an open chain. Note: rods require at least 2 segments. label: Optional label prefix for bodies, shapes, and joints. @@ -7573,7 +7730,7 @@ def add_rod( Note: - Bend defaults are 0.0 (no bending resistance unless specified). Stretch defaults to 1.0e5; pass a larger value when neighboring capsules should remain nearly inextensible. - - Stretch, bend, and damping values are passed through as provided per joint. + - Stretch, shear, bend, twist, and damping values are passed through as provided per joint. - Each segment is implemented as a capsule primitive. ``half_height`` is the half-length of the cylindrical centerline, excluding the hemispherical caps. - With ``body_frame_origin="start"``, the body origin is at the first centerline endpoint, @@ -7597,6 +7754,10 @@ def add_rod( # Input validation if stretch_stiffness < 0.0 or bend_stiffness < 0.0: raise ValueError("add_rod: stretch_stiffness and bend_stiffness must be >= 0") + if shear_stiffness is not None and shear_stiffness < 0.0: + raise ValueError("add_rod: shear_stiffness must be >= 0") + if twist_stiffness is not None and twist_stiffness < 0.0: + raise ValueError("add_rod: twist_stiffness must be >= 0") body_frame_origin = self._resolve_rod_body_frame_origin("add_rod", body_frame_origin) num_segments = len(positions) - 1 @@ -7636,8 +7797,12 @@ def add_rod( cfg=cfg, stretch_stiffness=stretch_stiffness, stretch_damping=stretch_damping, + shear_stiffness=shear_stiffness, + shear_damping=shear_damping, bend_stiffness=bend_stiffness, bend_damping=bend_damping, + twist_stiffness=twist_stiffness, + twist_damping=twist_damping, label=label, wrap_in_articulation=False, quaternions=quaternions, @@ -7691,8 +7856,12 @@ def add_rod( child_xform=child_xform, bend_stiffness=bend_stiffness, bend_damping=bend_damping, + twist_stiffness=twist_stiffness, + twist_damping=twist_damping, stretch_stiffness=stretch_stiffness, stretch_damping=stretch_damping, + shear_stiffness=shear_stiffness, + shear_damping=shear_damping, label=loop_joint_label, collision_filter_parent=True, enabled=True, @@ -7711,8 +7880,12 @@ def add_rod_graph( cfg: ShapeConfig | None = None, stretch_stiffness: float | None = None, stretch_damping: float | None = None, + shear_stiffness: float | None = None, + shear_damping: float | None = None, bend_stiffness: float | None = None, bend_damping: float | None = None, + twist_stiffness: float | None = None, + twist_damping: float | None = None, label: str | None = None, wrap_in_articulation: bool = True, quaternions: list[Quat] | None = None, @@ -7752,10 +7925,17 @@ def add_rod_graph( stretch_stiffness: Per-joint cable stretch stiffness, stored directly as ``target_ke`` [N/m]. Defaults to 1.0e5. stretch_damping: Stretch damping [N·s/m] (per joint). Defaults to 0.0. - bend_stiffness: Per-joint cable bend/twist stiffness, stored directly as ``target_ke`` - (torque per radian). + shear_stiffness: Optional per-joint transverse shear stiffness [N/m]. If None, defaults to + ``stretch_stiffness``. + shear_damping: Optional per-joint transverse shear damping [N·s/m]. If None, defaults to + ``stretch_damping`` only when both ``shear_stiffness`` and ``shear_damping`` are None. Otherwise defaults to 0.0. + bend_stiffness: Per-joint cable bend stiffness, stored directly as ``target_ke`` [N*m]. Defaults to 0.0. - bend_damping: Bend/twist damping [N·m·s/rad] (per joint). Defaults to 0.0. + bend_damping: Bend damping [N·m·s/rad] (per joint). Defaults to 0.0. + twist_stiffness: Optional per-joint cable twist stiffness [N*m]. If None, defaults to + ``bend_stiffness``. + twist_damping: Optional per-joint cable twist damping [N·m·s/rad]. If None, defaults to ``bend_damping`` + only when both ``twist_stiffness`` and ``twist_damping`` are None. Otherwise defaults to 0.0. label: Optional label prefix for bodies, shapes, joints, and articulations. wrap_in_articulation: If True, wraps the generated joint forest into one articulation per connected component. @@ -7797,6 +7977,10 @@ def add_rod_graph( if stretch_stiffness < 0.0 or bend_stiffness < 0.0: raise ValueError("add_rod_graph: stretch_stiffness and bend_stiffness must be >= 0") + if shear_stiffness is not None and shear_stiffness < 0.0: + raise ValueError("add_rod_graph: shear_stiffness must be >= 0") + if twist_stiffness is not None and twist_stiffness < 0.0: + raise ValueError("add_rod_graph: twist_stiffness must be >= 0") body_frame_origin = self._resolve_rod_body_frame_origin("add_rod_graph", body_frame_origin) if len(node_positions) < 2: raise ValueError("add_rod_graph: node_positions must contain at least 2 nodes") @@ -7958,8 +8142,12 @@ def _build_joints_star() -> list[int]: child_xform=child_xform, bend_stiffness=bend_stiffness, bend_damping=bend_damping, + twist_stiffness=twist_stiffness, + twist_damping=twist_damping, stretch_stiffness=stretch_stiffness, stretch_damping=stretch_damping, + shear_stiffness=shear_stiffness, + shear_damping=shear_damping, label=joint_label, collision_filter_parent=True, enabled=True, @@ -8013,8 +8201,12 @@ def _build_joints_forest() -> list[int]: child_xform=child_xform, bend_stiffness=bend_stiffness, bend_damping=bend_damping, + twist_stiffness=twist_stiffness, + twist_damping=twist_damping, stretch_stiffness=stretch_stiffness, stretch_damping=stretch_damping, + shear_stiffness=shear_stiffness, + shear_damping=shear_damping, label=joint_label, collision_filter_parent=True, enabled=True, @@ -8166,9 +8358,32 @@ def add_particles( Note: Set the mass equal to zero to create a 'kinematic' particle that is not subject to dynamics. + + Raises: + ValueError: If a particle input or custom attribute list has a mismatched length. """ - particle_start = self.particle_count particle_count = len(pos) + required_inputs = ( + ("vel", vel), + ("mass", mass), + ) + optional_inputs = ( + ("radius", radius), + ("flags", flags), + ) + for name, values in required_inputs: + if values is None or len(values) != particle_count: + actual_count = "None" if values is None else len(values) + raise ValueError( + f"{name} length mismatch: expected {particle_count} values to match pos, got {actual_count}" + ) + for name, values in optional_inputs: + if values is not None and len(values) != particle_count: + raise ValueError( + f"{name} length mismatch: expected {particle_count} values to match pos, got {len(values)}" + ) + + particle_start = self.particle_count self.particle_q.extend(pos) self.particle_qd.extend(vel) @@ -10407,7 +10622,21 @@ def _eq_index_array(name: str) -> np.ndarray: f"{next_start[idx]}." ) - # Validate array length consistency + particle_count = self.particle_count + particle_arrays = ( + ("particle_qd", self.particle_qd), + ("particle_mass", self.particle_mass), + ("particle_radius", self.particle_radius), + ("particle_flags", self.particle_flags), + ("particle_world", self.particle_world), + ) + for name, values in particle_arrays: + if len(values) != particle_count: + raise ValueError( + f"Array length mismatch: {name} has length {len(values)}, " + f"but expected {particle_count} (particle_count)." + ) + if joint_count > 0: # Per-DOF arrays should have length == joint_dof_count dof_arrays = [ @@ -11155,28 +11384,6 @@ def compute_voxel_resolution_from_aabb(aabb_lower, aabb_upper, voxel_budget): # --------------------- # Compute and compact texture SDF resources (shared table + per-shape index indirection) - from ..geometry.types import Mesh as NewtonMesh # noqa: PLC0415 - - def _create_primitive_mesh(stype: int, scale: Sequence[float] | None) -> NewtonMesh | None: - """Create a watertight mesh from a primitive shape for texture SDF construction.""" - from ..core.types import Axis # noqa: PLC0415 - - sx, sy, sz = scale if scale is not None else (1.0, 1.0, 1.0) - common_kw = {"compute_normals": False, "compute_uvs": False, "compute_inertia": False} - if stype == GeoType.BOX: - return NewtonMesh.create_box(sx, sy, sz, duplicate_vertices=False, **common_kw) - elif stype == GeoType.SPHERE: - return NewtonMesh.create_sphere(sx, **common_kw) - elif stype == GeoType.CAPSULE: - return NewtonMesh.create_capsule(sx, sy, up_axis=Axis.Z, **common_kw) - elif stype == GeoType.CYLINDER: - return NewtonMesh.create_cylinder(sx, sy, up_axis=Axis.Z, **common_kw) - elif stype == GeoType.CONE: - return NewtonMesh.create_cone(sx, sy, up_axis=Axis.Z, **common_kw) - elif stype == GeoType.ELLIPSOID: - return NewtonMesh.create_ellipsoid(sx, sy, sz, **common_kw) - return None - current_device = wp.get_device(device) is_gpu = current_device.is_cuda @@ -11220,6 +11427,7 @@ def _create_primitive_mesh(stype: int, scale: Sequence[float] | None) -> NewtonM TextureSDFData, create_empty_texture_sdf_data, create_texture_sdf_from_mesh, + create_texture_sdf_from_primitive, ) _tex_fmt_map = { @@ -11338,44 +11546,39 @@ def _create_primitive_mesh(stype: int, scale: Sequence[float] | None) -> NewtonM compact_texture_sdf_subgrid_textures.append(None) compact_texture_sdf_subgrid_start_slots.append(None) else: - prim_mesh = _create_primitive_mesh(shape_type, shape_scale) - if prim_mesh is not None: - prim_wp_mesh = wp.Mesh( - points=wp.array(prim_mesh.vertices, dtype=wp.vec3, device=device), - indices=wp.array(prim_mesh.indices.flatten(), dtype=wp.int32, device=device), - support_winding_number=True, - ) - try: - tex_data, c_tex, s_tex = create_texture_sdf_from_mesh( - prim_wp_mesh, - margin=sdf_gen_margin, - narrow_band_range=tuple(sdf_narrow_band_range), - max_resolution=effective_max_resolution, - target_voxel_size=sdf_target_voxel_size, - quantization_mode=_tex_fmt_map[sdf_tex_fmt], - scale_baked=True, - device=device, - ) - except Exception as e: - warnings.warn( - f"Texture SDF construction failed for shape {i} " - f"(type={shape_type}): {e}. Falling back to BVH.", - stacklevel=2, - ) - tex_data = create_empty_texture_sdf_data() - c_tex = None - s_tex = None - compact_texture_sdf_data.append(tex_data) - compact_texture_sdf_coarse_textures.append(c_tex) - compact_texture_sdf_subgrid_textures.append(s_tex) - compact_texture_sdf_subgrid_start_slots.append( - tex_data.subgrid_start_slots if c_tex is not None else None + try: + tex_data, c_tex, s_tex = create_texture_sdf_from_primitive( + shape_type, + shape_scale, + margin=sdf_gen_margin, + narrow_band_range=tuple(sdf_narrow_band_range), + max_resolution=effective_max_resolution, + target_voxel_size=sdf_target_voxel_size, + quantization_mode=_tex_fmt_map[sdf_tex_fmt], + scale_baked=True, + device=device, ) - else: + except NotImplementedError: compact_texture_sdf_data.append(create_empty_texture_sdf_data()) compact_texture_sdf_coarse_textures.append(None) compact_texture_sdf_subgrid_textures.append(None) compact_texture_sdf_subgrid_start_slots.append(None) + continue + except Exception as e: + warnings.warn( + f"Texture SDF construction failed for shape {i} " + f"(type={shape_type}): {e}. Falling back to BVH.", + stacklevel=2, + ) + tex_data = create_empty_texture_sdf_data() + c_tex = None + s_tex = None + compact_texture_sdf_data.append(tex_data) + compact_texture_sdf_coarse_textures.append(c_tex) + compact_texture_sdf_subgrid_textures.append(s_tex) + compact_texture_sdf_subgrid_start_slots.append( + tex_data.subgrid_start_slots if c_tex is not None else None + ) # Build volume SDFs for participating MESH/CONVEX_MESH shapes that still lack one, when a # per-shape SDF is requested -- ShapeConfig.configure_sdf(force_sdf=True), or an sdf @@ -11964,7 +12167,11 @@ def _to_wp_array(data, dtype, requires_grad): # Add custom attributes onto the model (with lazy evaluation) # Early return if no custom attributes exist to avoid overhead if not self.custom_attributes: - m.bvh_build_shapes(m, bvh_constructor=self.default_bvh_cfg.shape_constructor) + m.bvh_build_shapes( + m, + bvh_constructor=self.default_bvh_cfg.shape_constructor, + shape_flags=self.default_bvh_cfg.shape_flags, + ) m.bvh_build_particles(m) return m @@ -12071,7 +12278,11 @@ def _to_wp_array(data, dtype, requires_grad): custom_attr.references, ) - m.bvh_build_shapes(m, bvh_constructor=self.default_bvh_cfg.shape_constructor) + m.bvh_build_shapes( + m, + bvh_constructor=self.default_bvh_cfg.shape_constructor, + shape_flags=self.default_bvh_cfg.shape_flags, + ) m.bvh_build_particles(m) return m diff --git a/newton/_src/sim/collide.py b/newton/_src/sim/collide.py index af35044a5c..07e971fecd 100644 --- a/newton/_src/sim/collide.py +++ b/newton/_src/sim/collide.py @@ -562,6 +562,28 @@ def _build_soft_particle_rigid_contact_pairs(model: Model) -> wp.array[wp.vec2i] return _world_compatible_pairs(model.particle_world.numpy(), model.shape_world.numpy(), world_count, model.device) +def _count_soft_particle_rigid_contact_pairs(model: Model) -> int: + """Count exactly how many pairs :func:`_build_soft_particle_rigid_contact_pairs` emits for ``model``. + + Reads only the per-world start offsets, so solvers can pre-size soft-contact buffers without + downloading per-entity world ids. This is not :attr:`CollisionPipeline.soft_contact_max`, which + additionally reserves edge/face headroom when ``enable_rigid_soft_full_surface_contact`` is set. + Reads host arrays, so it is not graph-capture-safe; call at solver construction. + """ + particle_start = model.particle_world_start.numpy() + shape_start = model.shape_world_start.numpy() + global_particles = int(particle_start[-1] - particle_start[-2] + particle_start[0]) + global_shapes = int(shape_start[-1] - shape_start[-2] + shape_start[0]) + # Global particles pair with every shape; local particles additionally pair with global shapes. + total = global_particles * model.shape_count + total += (model.particle_count - global_particles) * global_shapes + # Local particles pair with the shapes sharing their world. + per_world = slice(0, model.world_count + 1) + return total + int( + np.dot(np.diff(particle_start[per_world]).astype(np.int64), np.diff(shape_start[per_world]).astype(np.int64)) + ) + + def _build_soft_face_rigid_contact_pairs( model: Model, capable_shape_mask: np.ndarray | None = None ) -> wp.array[wp.vec2i]: @@ -780,7 +802,8 @@ def __init__( soft_contact_max: Maximum number of soft contacts to allocate. If None, defaults to ``soft_rigid_contact_pair_count``, the number of precomputed soft-rigid (particle-shape) pairs launched for soft - contact generation. + contact generation, plus the full-surface edge/face headroom when + ``enable_rigid_soft_full_surface_contact`` is set. soft_contact_margin: Margin for soft contact generation. Defaults to 0.01. enable_rigid_soft_full_surface_contact: Generate soft contacts over the full soft-mesh surface -- the edges and triangle interiors -- against rigid SDFs, in addition to the @@ -1179,7 +1202,8 @@ def soft_contact_max(self) -> int: def soft_rigid_contact_pair_count(self) -> int: """Number of precomputed soft-rigid (particle-shape) pairs launched for soft contacts. - This is the default capacity used for ``soft_contact_max``. + This is the base of the default ``soft_contact_max``, which additionally reserves + edge/face headroom when ``enable_rigid_soft_full_surface_contact`` is set. """ return self._soft_rigid_contact_pair_count diff --git a/newton/_src/sim/ik/ik_lm_optimizer.py b/newton/_src/sim/ik/ik_lm_optimizer.py index a7e688f99f..b43befef4e 100644 --- a/newton/_src/sim/ik/ik_lm_optimizer.py +++ b/newton/_src/sim/ik/ik_lm_optimizer.py @@ -80,6 +80,44 @@ def _update_lm_state( lambda_values[row] = wp.clamp(new_lambda, lambda_min, lambda_max) +@wp.kernel +def _zero_fixed_dof_jacobian_columns( + joint_dof_mask: wp.array[wp.bool], + jacobian: wp.array3d[wp.float32], +): + row, residual, dof = wp.tid() + if not joint_dof_mask[dof]: + jacobian[row, residual, dof] = 0.0 + + +def _validate_joint_dof_mask(model: Model, joint_dof_mask: wp.array[wp.bool]) -> None: + if joint_dof_mask.dtype != wp.bool: + raise ValueError("joint_dof_mask must have dtype wp.bool") + if joint_dof_mask.ndim != 1 or joint_dof_mask.shape[0] != model.joint_dof_count: + raise ValueError("joint_dof_mask must have shape [joint_dof_count]") + if joint_dof_mask.device != model.device: + raise ValueError("joint_dof_mask must be on the model device") + + # The mask acts on twist-space DOFs, but the integrator couples a + # quaternion-integrated joint's DOFs to its coordinates (rotation about the + # joint origin translates the body), so a partial mask would not keep the + # remaining coordinates fixed. Require all-or-nothing masks for such joints. + mask = joint_dof_mask.numpy() + joint_type = model.joint_type.numpy() + qd_start = model.joint_qd_start.numpy() + quaternion_joints = (JointType.BALL, JointType.FREE, JointType.DISTANCE) + for j in range(len(joint_type)): + if joint_type[j] not in quaternion_joints: + continue + joint_mask = mask[qd_start[j] : qd_start[j + 1]] + if joint_mask.any() and not joint_mask.all(): + raise ValueError( + f"joint_dof_mask partially masks joint {j} " + f"({JointType(joint_type[j]).name}): quaternion-integrated joints " + "must have all of their DOFs masked together" + ) + + class IKOptimizerLM: """Levenberg-Marquardt optimizer for batched inverse kinematics. @@ -103,6 +141,13 @@ class IKOptimizerLM: accept a step. problem_idx: Optional mapping from batch rows to base problem indices for per-problem objective data. + joint_dof_mask: Optional model-wide mask, shape ``[joint_dof_count]``, + indexed in DOF (velocity) space per :attr:`Model.joint_qd_start` — + a free joint has 6 entries. ``True`` entries are optimized; + ``False`` entries receive an exactly-zero update. Quaternion- + integrated joints (free/ball/distance) must be masked + all-or-nothing, which the constructor enforces. The mask array must + not be modified after construction. """ TILE_N_DOFS = None @@ -142,6 +187,7 @@ def __init__( rho_min: float = 1e-3, *, problem_idx: wp.array[wp.int32] | None = None, + joint_dof_mask: wp.array[wp.bool] | None = None, ) -> None: self.model = model self.device = model.device @@ -160,6 +206,9 @@ def __init__( self.lambda_min = lambda_min self.lambda_max = lambda_max self.rho_min = rho_min + if joint_dof_mask is not None: + _validate_joint_dof_mask(model, joint_dof_mask) + self.joint_dof_mask = joint_dof_mask if self.TILE_N_DOFS is not None: assert self.n_dofs == self.TILE_N_DOFS @@ -365,10 +414,12 @@ def _jacobian_at(self, ctx: BatchCtx) -> wp.array3d[wp.float32]: if mode == IKJacobianType.AUTODIFF: self._jacobian_autodiff(ctx) + self._apply_joint_dof_mask(ctx.jacobian_out) return ctx.jacobian_out if mode == IKJacobianType.ANALYTIC: self._jacobian_analytic(ctx, accumulate=False) + self._apply_joint_dof_mask(ctx.jacobian_out) return ctx.jacobian_out # MIXED mode @@ -380,8 +431,20 @@ def _jacobian_at(self, ctx: BatchCtx) -> wp.array3d[wp.float32]: if self.has_analytic_objective: self._jacobian_analytic(ctx, accumulate=self.has_autodiff_objective) + self._apply_joint_dof_mask(ctx.jacobian_out) return ctx.jacobian_out + def _apply_joint_dof_mask(self, jacobian: wp.array3d[wp.float32]) -> None: + if self.joint_dof_mask is None: + return + wp.launch( + _zero_fixed_dof_jacobian_columns, + dim=(self.n_batch, self.n_residuals, self.n_dofs), + inputs=[self.joint_dof_mask], + outputs=[jacobian], + device=self.device, + ) + def _jacobian_autodiff(self, ctx: BatchCtx) -> None: if self.tape is None: raise RuntimeError("Autodiff Jacobian requested but tape is not initialized") diff --git a/newton/_src/sim/ik/ik_solver.py b/newton/_src/sim/ik/ik_solver.py index 1746f6a6d0..f8d30ca46a 100644 --- a/newton/_src/sim/ik/ik_solver.py +++ b/newton/_src/sim/ik/ik_solver.py @@ -217,6 +217,14 @@ class IKSolver: lambda_min: Minimum LM damping value. lambda_max: Maximum LM damping value. rho_min: Minimum LM acceptance ratio. + joint_dof_mask: Optional model-wide mask, shape ``[joint_dof_count]``, + indexed in DOF (velocity) space per :attr:`Model.joint_qd_start` — + a free joint has 6 entries. ``True`` entries are optimized; + ``False`` entries receive an exactly-zero update. Quaternion- + integrated joints (free/ball/distance) must be masked + all-or-nothing, which the constructor enforces. The mask array + must not be modified after construction. Currently supported by + the LM optimizer with sampling disabled. history_len: Number of correction pairs retained by L-BFGS. h0_scale: Initial inverse-Hessian scale for L-BFGS. line_search_alphas: Candidate line-search step sizes for L-BFGS. @@ -242,6 +250,7 @@ def __init__( lambda_min: float = 1e-5, lambda_max: float = 1e10, rho_min: float = 1e-3, + joint_dof_mask: wp.array[wp.bool] | None = None, # L-BFGS parameters history_len: int = 10, h0_scale: float = 1.0, @@ -260,6 +269,12 @@ def __init__( raise ValueError("n_seeds must be >= 1") if sampler is IKSampler.NONE and n_seeds != 1: raise ValueError("sampler 'none' requires n_seeds == 1") + if joint_dof_mask is not None: + if optimizer is not IKOptimizer.LM: + raise ValueError("joint_dof_mask is only supported by the LM optimizer") + if sampler is not IKSampler.NONE: + raise ValueError("joint_dof_mask requires sampler='none'") + # remaining mask validation happens in IKOptimizerLM self.model = model self.device = model.device @@ -310,6 +325,7 @@ def __init__( lambda_min=lambda_min, lambda_max=lambda_max, rho_min=rho_min, + joint_dof_mask=joint_dof_mask, ) elif optimizer is IKOptimizer.LBFGS: self._impl = IKOptimizerLBFGS( diff --git a/newton/_src/sim/model.py b/newton/_src/sim/model.py index 4da7bc8dd6..b6864121e4 100644 --- a/newton/_src/sim/model.py +++ b/newton/_src/sim/model.py @@ -18,6 +18,7 @@ import warp as wp from ..core.types import Devicelike, override +from ..geometry.flags import ShapeFlags from ..utils.mesh import MeshAdjacency, MeshAdjacencyData from .contacts import Contacts from .control import Control @@ -31,11 +32,6 @@ from .collide import CollisionPipeline -_HAS_HEIGHTFIELDS_DEPRECATION_MSG = ( - "Model.has_heightfields is deprecated; use Model.heightfield_count, " - "or model.heightfield_count > 0 for boolean checks, instead." -) - _SHAPE_COLLISION_FILTER_MUTATION_DEPRECATION_MSG = ( "Mutating Model.shape_collision_filter_pairs after ModelBuilder.finalize() is deprecated. " "Configure collision filters on ModelBuilder before finalizing; post-finalize filter changes " @@ -900,7 +896,7 @@ def __init__(self, device: Devicelike | None = None): # Shape and particle BVH structures and related fields self.bvh_shapes: wp.Bvh | None = None - """BVH over visible shapes, indexed by ``bvh_shape_enabled``. Built by :meth:`ModelBuilder.finalize`.""" + """BVH over selected shapes, indexed by ``bvh_shape_enabled``. Built by :meth:`ModelBuilder.finalize`.""" self.bvh_shapes_group_roots: wp.array[wp.int32] | None = None """Per-world BVH group roots for shapes, shape ``[world_count + 1]`` (last slot is global).""" self.bvh_shape_enabled: wp.array[wp.uint32] | None = None @@ -957,11 +953,6 @@ def __init__(self, device: Devicelike | None = None): self._texture_sdf_subgrid_start_slots: list = [] """Subgrid start slot arrays matching _texture_sdf_data by index. Kept for reference counting.""" - # Caches for the deprecated lazy ``sdf_block_coords`` / ``sdf_index2blocks`` - # properties. Populated on first access; cleared when SDF storage changes. - self._sdf_block_coords_cache: wp.array | None = None - self._sdf_index2blocks_cache: wp.array | None = None - # Local AABB and voxel grid for contact reduction # Note: These are stored in Model (not Contacts) because they are static geometry properties # computed once during finalization, not per-frame contact data. @@ -1568,240 +1559,6 @@ def _normalize_attribute_reference(self, references: str | None) -> Model.Attrib return references raise ValueError(f"Unknown custom attribute reference frequency {references!r}") - # ----- Deprecated SDF aliases ------------------------------------------- - # The underlying SDF members on ``Model`` are now underscore-prefixed. - # The properties below preserve the historical attribute names for one - # release cycle and emit ``DeprecationWarning`` on access. - - @property - def shape_sdf_index(self) -> wp.array[wp.int32] | None: - """Deprecated alias for :attr:`_shape_sdf_index`. - - .. deprecated:: 1.3 - Use the underscored private member or the appropriate accessor. - This alias will be removed in a future release. - """ - warnings.warn( - "Model.shape_sdf_index is deprecated; use Model._shape_sdf_index. " - "The public alias will be removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - return self._shape_sdf_index - - @shape_sdf_index.setter - def shape_sdf_index(self, value): - warnings.warn( - "Model.shape_sdf_index is deprecated; assign to Model._shape_sdf_index. " - "The public alias will be removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - self._shape_sdf_index = value - - @property - def texture_sdf_data(self): - """Deprecated alias for :attr:`_texture_sdf_data`. - - .. deprecated:: 1.3 - Use the underscored private member. The alias will be removed in - a future release. - """ - warnings.warn( - "Model.texture_sdf_data is deprecated; use Model._texture_sdf_data. " - "The public alias will be removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - return self._texture_sdf_data - - @texture_sdf_data.setter - def texture_sdf_data(self, value): - warnings.warn( - "Model.texture_sdf_data is deprecated; assign to Model._texture_sdf_data. " - "The public alias will be removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - self._texture_sdf_data = value - self._sdf_block_coords_cache = None - self._sdf_index2blocks_cache = None - - @property - def texture_sdf_coarse_textures(self) -> list: - """Deprecated alias for :attr:`_texture_sdf_coarse_textures`. - - .. deprecated:: 1.3 - Use the underscored private member. The alias will be removed in - a future release. - """ - warnings.warn( - "Model.texture_sdf_coarse_textures is deprecated; use " - "Model._texture_sdf_coarse_textures. The public alias will be " - "removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - return self._texture_sdf_coarse_textures - - @texture_sdf_coarse_textures.setter - def texture_sdf_coarse_textures(self, value): - warnings.warn( - "Model.texture_sdf_coarse_textures is deprecated; assign to " - "Model._texture_sdf_coarse_textures. The public alias will be " - "removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - self._texture_sdf_coarse_textures = value - self._sdf_block_coords_cache = None - self._sdf_index2blocks_cache = None - - @property - def texture_sdf_subgrid_textures(self) -> list: - """Deprecated alias for :attr:`_texture_sdf_subgrid_textures`. - - .. deprecated:: 1.3 - Use the underscored private member. The alias will be removed in - a future release. - """ - warnings.warn( - "Model.texture_sdf_subgrid_textures is deprecated; use " - "Model._texture_sdf_subgrid_textures. The public alias will be " - "removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - return self._texture_sdf_subgrid_textures - - @texture_sdf_subgrid_textures.setter - def texture_sdf_subgrid_textures(self, value): - warnings.warn( - "Model.texture_sdf_subgrid_textures is deprecated; assign to " - "Model._texture_sdf_subgrid_textures. The public alias will be " - "removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - self._texture_sdf_subgrid_textures = value - - @property - def texture_sdf_subgrid_start_slots(self) -> list: - """Deprecated alias for :attr:`_texture_sdf_subgrid_start_slots`. - - .. deprecated:: 1.3 - Use the underscored private member. The alias will be removed in - a future release. - """ - warnings.warn( - "Model.texture_sdf_subgrid_start_slots is deprecated; use " - "Model._texture_sdf_subgrid_start_slots. The public alias will be " - "removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - return self._texture_sdf_subgrid_start_slots - - @texture_sdf_subgrid_start_slots.setter - def texture_sdf_subgrid_start_slots(self, value): - warnings.warn( - "Model.texture_sdf_subgrid_start_slots is deprecated; assign to " - "Model._texture_sdf_subgrid_start_slots. The public alias will be " - "removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - self._texture_sdf_subgrid_start_slots = value - - @property - def sdf_block_coords(self): - """Deprecated. Lazily-computed flat ``wp.vec3us`` block coords. - - Per-SDF active-block coordinates were dropped when the hydroelastic - broadphase started deriving them arithmetically from each SDF's - coarse-texture dimensions. This property recomputes the legacy - layout on first access (and caches it) so external callers that - still read the attribute keep working. - - .. deprecated:: 1.3 - This attribute will be removed in a future release. - """ - warnings.warn( - "Model.sdf_block_coords is deprecated and will be removed in " - "a future release. The hydroelastic broadphase now derives block " - "coordinates arithmetically from each SDF's coarse-texture " - "dimensions and no longer needs this attribute.", - DeprecationWarning, - stacklevel=2, - ) - self._ensure_legacy_sdf_block_arrays() - return self._sdf_block_coords_cache - - @property - def sdf_index2blocks(self): - """Deprecated. Lazily-computed per-SDF ``[start, end)`` ranges. - - Per-SDF ``[start, end)`` indices into ``sdf_block_coords`` were - dropped when the hydroelastic broadphase started deriving block - ranges arithmetically from each SDF's coarse-texture dimensions. - This property recomputes the legacy layout on first access (and - caches it) so external callers that still read the attribute keep - working. - - .. deprecated:: 1.3 - This attribute will be removed in a future release. - """ - warnings.warn( - "Model.sdf_index2blocks is deprecated and will be removed in " - "a future release. The hydroelastic broadphase now derives block " - "ranges arithmetically from each SDF's coarse-texture " - "dimensions and no longer needs this attribute.", - DeprecationWarning, - stacklevel=2, - ) - self._ensure_legacy_sdf_block_arrays() - return self._sdf_index2blocks_cache - - def _ensure_legacy_sdf_block_arrays(self) -> None: - """Populate the legacy SDF block-coord caches on demand.""" - if self._sdf_block_coords_cache is not None and self._sdf_index2blocks_cache is not None: - return - # Local import keeps the deprecated module out of the normal load path. - from ..geometry._deprecated_sdf_block_coords import ( # noqa: PLC0415 - build_legacy_sdf_block_arrays, - ) - - subgrid_size = 8 - if self._texture_sdf_data is not None and len(self._texture_sdf_data) > 0: - subgrid_size = int(self._texture_sdf_data.numpy()[0]["subgrid_size"]) - block_coords, index2blocks = build_legacy_sdf_block_arrays( - self._texture_sdf_coarse_textures, - subgrid_size=subgrid_size, - device=self.device, - ) - self._sdf_block_coords_cache = block_coords - self._sdf_index2blocks_cache = index2blocks - - @property - def has_heightfields(self) -> bool: - """Deprecated boolean alias for :attr:`heightfield_count`. - - .. deprecated:: 1.3 - Use :attr:`heightfield_count`, or ``heightfield_count > 0`` for - boolean checks, instead. - """ - import warnings # noqa: PLC0415 - - warnings.warn(_HAS_HEIGHTFIELDS_DEPRECATION_MSG, DeprecationWarning, stacklevel=2) - return self.heightfield_count > 0 - - @has_heightfields.setter - def has_heightfields(self, value: bool) -> None: - import warnings # noqa: PLC0415 - - warnings.warn(_HAS_HEIGHTFIELDS_DEPRECATION_MSG, DeprecationWarning, stacklevel=2) - self.heightfield_count = 1 if value else 0 - @property def joint_target_q_start(self) -> wp.array | None: """Per-joint start index into :attr:`joint_target_q`, shape @@ -1885,7 +1642,13 @@ def joint_target_vel(self, value: wp.array | None) -> None: ) self.joint_target_qd = value - def bvh_build_shapes(self, state: State, *, bvh_constructor: str | None = None) -> None: + def bvh_build_shapes( + self, + state: State, + *, + bvh_constructor: str | None = None, + shape_flags: ShapeFlags = ShapeFlags.VISIBLE, + ) -> None: """Build or rebuild the shape BVH stored on this model. Allocates :attr:`bvh_shapes` and related fields from the current @@ -1899,6 +1662,8 @@ def bvh_build_shapes(self, state: State, *, bvh_constructor: str | None = None) bvh_constructor: Warp BVH construction algorithm. Valid choices are ``"sah"``, ``"median"``, ``"lbvh"``, or ``None`` to use Warp's device-dependent default. + shape_flags: Mask of :class:`~newton.ShapeFlags`; a shape is + included in the BVH if any of its flags are set in the mask. """ from ..geometry.bvh import ( # noqa: PLC0415 compute_bvh_group_roots, @@ -1936,6 +1701,7 @@ def bvh_build_shapes(self, state: State, *, bvh_constructor: str | None = None) inputs=[ self.shape_type, self.shape_flags, + int(shape_flags), self.bvh_shape_enabled, num_enabled, ], @@ -1945,6 +1711,9 @@ def bvh_build_shapes(self, state: State, *, bvh_constructor: str | None = None) self.bvh_shape_world_transforms = wp.empty(shape_count, dtype=wp.transformf, device=device) if self.bvh_shape_count_enabled == 0: + # drop any BVH from a previous build, it would index stale shapes + self.bvh_shapes = None + self.bvh_shapes_group_roots = None return compute_shape_world_transforms_launch(self, state) diff --git a/newton/_src/solvers/__init__.py b/newton/_src/solvers/__init__.py index 3f6a211e32..d6bcfd9852 100644 --- a/newton/_src/solvers/__init__.py +++ b/newton/_src/solvers/__init__.py @@ -1,16 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers # SPDX-License-Identifier: Apache-2.0 -from .featherstone import SolverFeatherstone -from .flags import SolverNotifyFlags -from .implicit_mpm import SolverImplicitMPM -from .kamino import SolverKamino -from .mujoco import SolverMuJoCo -from .semi_implicit import SolverSemiImplicit -from .solver import SolverBase -from .style3d.solver_style3d import SolverStyle3D -from .vbd import SolverVBD -from .xpbd import SolverXPBD +import importlib +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from . import style3d + from .featherstone import SolverFeatherstone + from .implicit_mpm import SolverImplicitMPM + from .kamino import SolverKamino + from .mujoco import SolverMuJoCo + from .semi_implicit import SolverSemiImplicit + from .solver import SolverBase + from .style3d.solver_style3d import SolverStyle3D + from .vbd import SolverVBD + from .xpbd import SolverXPBD __all__ = [ "SolverBase", @@ -18,9 +22,42 @@ "SolverImplicitMPM", "SolverKamino", "SolverMuJoCo", - "SolverNotifyFlags", "SolverSemiImplicit", "SolverStyle3D", "SolverVBD", "SolverXPBD", + "style3d", ] + +# Maps each public symbol to the module that provides it and the attribute to +# fetch from that module (None returns the module itself). Symbols are +# resolved on first attribute access (PEP 562) so that importing Newton does +# not pay the import cost of every solver backend. +_LAZY_IMPORTS: dict[str, tuple[str, str | None]] = { + "SolverBase": (".solver", "SolverBase"), + "SolverFeatherstone": (".featherstone", "SolverFeatherstone"), + "SolverImplicitMPM": (".implicit_mpm", "SolverImplicitMPM"), + "SolverKamino": (".kamino", "SolverKamino"), + "SolverMuJoCo": (".mujoco", "SolverMuJoCo"), + "SolverSemiImplicit": (".semi_implicit", "SolverSemiImplicit"), + "SolverStyle3D": (".style3d.solver_style3d", "SolverStyle3D"), + "SolverVBD": (".vbd", "SolverVBD"), + "SolverXPBD": (".xpbd", "SolverXPBD"), + "style3d": (".style3d", None), +} + + +def __getattr__(name: str): + try: + module_name, attr_name = _LAZY_IMPORTS[name] + except KeyError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + + module = importlib.import_module(module_name, __name__) + value = module if attr_name is None else getattr(module, attr_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(_LAZY_IMPORTS)) diff --git a/newton/_src/solvers/coupled/solver_coupled.py b/newton/_src/solvers/coupled/solver_coupled.py index 11ebad64da..63241cb7c0 100644 --- a/newton/_src/solvers/coupled/solver_coupled.py +++ b/newton/_src/solvers/coupled/solver_coupled.py @@ -14,7 +14,7 @@ import warp as wp from ...geometry import ParticleFlags, ShapeFlags -from ...sim import ModelFlags, StateFlags +from ...sim import Model, ModelFlags, StateFlags from ..solver import SolverBase from .interface import ( CouplingEndpointKind, @@ -23,7 +23,7 @@ from .model_view import ModelView, _AttributeNamespaceView if TYPE_CHECKING: - from ...sim import Contacts, Control, Model, State + from ...sim import Contacts, Control, State logger = logging.getLogger(__name__) @@ -1975,7 +1975,7 @@ def prepare_contacts(self, contacts: Contacts | None) -> None: def _sync_entry_reset_state(self, entry: SolverEntry) -> None: """Mirror a reset entry input state to persistent entry buffers.""" if entry.state_1 is not None and entry.state_1 is not entry.state_0: - _copy_state(entry.state_0, entry.state_1) + _copy_same_view_state(entry.state_0, entry.state_1) for entry_state in (entry.state_0, entry.state_1): if entry_state is not None: @@ -2283,7 +2283,7 @@ def _step_entry( substep_dt = dt / float(entry.substeps) if entry.state_tmp is None: raise RuntimeError(f"SolverCoupled.Entry {entry.name!r} is missing a substep scratch state") - _copy_state(entry.state_0, entry.state_1) + _copy_same_view_state(entry.state_0, entry.state_1) state_in = entry.state_1 state_out = entry.state_tmp for substep in range(entry.substeps): @@ -2292,7 +2292,7 @@ def _step_entry( entry.solver.step(state_in, state_out, control, contacts, substep_dt) state_in, state_out = state_out, state_in if state_in is entry.state_tmp: - _copy_state(entry.state_tmp, entry.state_1) + _copy_same_view_state(entry.state_tmp, entry.state_1) return contacts def _contacts_for_entry(self, entry: SolverEntry, contacts: Contacts | None) -> Contacts | None: @@ -2814,6 +2814,43 @@ def _copy_state(src: State, dst: State) -> None: _copy_prefix(dst.joint_qd, src.joint_qd, "joint_qd") +def _copy_same_view_state(src: State, dst: State) -> None: + """Copy persistent arrays between states allocated from the same model view. + + Core state arrays retain the coupled solver's established copy semantics, + while custom namespaced arrays (such as implicit MPM history) are copied in + full. Solver-produced top-level output annotations are intentionally not + copied; they are not persistent state and may only exist on output buffers. + """ + if src is dst: + return + + _copy_state(src, dst) + + src_namespaces = {name: value for name, value in vars(src).items() if isinstance(value, Model.AttributeNamespace)} + dst_namespaces = {name: value for name, value in vars(dst).items() if isinstance(value, Model.AttributeNamespace)} + for namespace_name in src_namespaces.keys() | dst_namespaces.keys(): + src_namespace = src_namespaces.get(namespace_name) + dst_namespace = dst_namespaces.get(namespace_name) + src_arrays = ( + {} + if src_namespace is None + else {name: value for name, value in vars(src_namespace).items() if isinstance(value, wp.array)} + ) + dst_arrays = ( + {} + if dst_namespace is None + else {name: value for name, value in vars(dst_namespace).items() if isinstance(value, wp.array)} + ) + for array_name in src_arrays.keys() | dst_arrays.keys(): + qualified_name = f"{namespace_name}.{array_name}" + if array_name not in dst_arrays: + raise ValueError(f"State is missing array for '{qualified_name}' which is present in the other state.") + if array_name not in src_arrays: + raise ValueError(f"Other state is missing array for '{qualified_name}' which is present in this state.") + dst_arrays[array_name].assign(src_arrays[array_name]) + + def _copy_forces(src: State, dst: State) -> None: """Copy force buffers without disturbing positions or velocities.""" if dst.body_f is not None: diff --git a/newton/_src/solvers/coupled/solver_coupled_proxy.py b/newton/_src/solvers/coupled/solver_coupled_proxy.py index 68fb3d211b..ef0bc29d69 100644 --- a/newton/_src/solvers/coupled/solver_coupled_proxy.py +++ b/newton/_src/solvers/coupled/solver_coupled_proxy.py @@ -112,6 +112,53 @@ class _ProxyRelaxationMode(IntEnum): } +@wp.func +def _reset_world_selected(world: int, world_mask: wp.array[wp.bool]) -> bool: + if world >= 0 and world < world_mask.shape[0]: + return world_mask[world] + return world < 0 and world_mask.shape[0] == 1 and world_mask[0] + + +@wp.kernel(module="unique", enable_backward=False) +def _zero_global_proxy_values_masked_kernel( + proxy_ids_global: wp.array[int], + entity_world: wp.array[int], + world_mask: wp.array[wp.bool], + values: wp.array[Any], +): + index = wp.tid() + proxy_id = proxy_ids_global[index] + if _reset_world_selected(entity_world[proxy_id], world_mask): + values[proxy_id] = values.dtype(0.0) + + +@wp.kernel(module="unique", enable_backward=False) +def _zero_proxy_row_values_masked_kernel( + proxy_ids_global: wp.array[int], + entity_world: wp.array[int], + world_mask: wp.array[wp.bool], + values: wp.array[Any], +): + index = wp.tid() + proxy_id = proxy_ids_global[index] + if _reset_world_selected(entity_world[proxy_id], world_mask): + values[index] = values.dtype(0.0) + + +@wp.kernel(module="unique", enable_backward=False) +def _zero_local_proxy_values_masked_kernel( + proxy_ids_global: wp.array[int], + proxy_ids_local: wp.array[int], + entity_world: wp.array[int], + world_mask: wp.array[wp.bool], + values: wp.array[Any], +): + index = wp.tid() + proxy_id = proxy_ids_global[index] + if _reset_world_selected(entity_world[proxy_id], world_mask): + values[proxy_ids_local[index]] = values.dtype(0.0) + + @wp.kernel(enable_backward=False) def _copy_indexed_float_kernel( src_indices: wp.array[int], @@ -408,19 +455,24 @@ def _validate_proxy_destination_ids_not_owned( ) @staticmethod - def _validate_proxy_body_worlds(model: Model, source_ids: Sequence[int], proxy_ids: Sequence[int]) -> None: - if model.body_world is None: + def _validate_proxy_entity_worlds( + entity_world: wp.array[int] | None, + source_ids: Sequence[int], + proxy_ids: Sequence[int], + entity_name: str, + ) -> None: + if entity_world is None: return - body_world = model.body_world.numpy() + worlds = entity_world.numpy() for source_id, proxy_id in zip(source_ids, proxy_ids, strict=True): - source_world = int(body_world[source_id]) - proxy_world = int(body_world[proxy_id]) + source_world = int(worlds[source_id]) + proxy_world = int(worlds[proxy_id]) if source_world != proxy_world: raise ValueError( - "Proxy source body and destination proxy body must live in the same world: " - f"source body {source_id} is in world {source_world}, " - f"proxy body {proxy_id} is in world {proxy_world}" + f"Proxy source {entity_name} and destination proxy {entity_name} must live in the same world: " + f"source {entity_name} {source_id} is in world {source_world}, " + f"proxy {entity_name} {proxy_id} is in world {proxy_world}" ) def _proxy_body_sets_by_destination(self) -> dict[str, set[int]]: @@ -602,8 +654,8 @@ def _build_proxy_entity_mappings( self._validate_proxy_ids(f"Proxy destination {entity_name}", proxy_local_ids, entity_count) self._validate_unique_proxy_ids(f"source {entity_name}", src_ids) self._validate_unique_proxy_ids(f"proxy {entity_name}", proxy_local_ids) - if is_body: - self._validate_proxy_body_worlds(model, src_ids, proxy_local_ids) + entity_world = model.body_world if is_body else model.particle_world + self._validate_proxy_entity_worlds(entity_world, src_ids, proxy_local_ids, entity_name) self._validate_proxy_source_ids_owned( entity_name, src_ids, @@ -1032,21 +1084,55 @@ def _reset_coupling_state( ) -> None: """Clear lagged proxy feedback and collision caches after reset.""" super()._reset_coupling_state(state, world_mask=world_mask, flags=flags) - for mapping in [*self._proxy_mappings, *self._proxy_particle_mappings]: + for mapping, entity_world in ( + *((mapping, self.model.body_world) for mapping in self._proxy_mappings), + *((mapping, self.model.particle_world) for mapping in self._proxy_particle_mappings), + ): + if world_mask is None or entity_world is None: + if mapping.coupling_forces is not None: + mapping.coupling_forces.zero_() + if mapping.coupling_forces_previous is not None: + mapping.coupling_forces_previous.zero_() + if mapping.aitken_residual_previous is not None: + mapping.aitken_residual_previous.zero_() + if mapping.aitken_stats is not None: + mapping.aitken_stats.zero_() + if mapping.aitken_relaxation is not None: + mapping.aitken_relaxation.fill_(mapping.proxy_relaxation) + if mapping.aitken_has_previous is not None: + mapping.aitken_has_previous.zero_() + if mapping.proxy_qd_before is not None: + mapping.proxy_qd_before.zero_() + continue + if mapping.coupling_forces is not None: - mapping.coupling_forces.zero_() - if mapping.coupling_forces_previous is not None: - mapping.coupling_forces_previous.zero_() - if mapping.aitken_residual_previous is not None: - mapping.aitken_residual_previous.zero_() - if mapping.aitken_stats is not None: - mapping.aitken_stats.zero_() - if mapping.aitken_relaxation is not None: - mapping.aitken_relaxation.fill_(mapping.proxy_relaxation) - if mapping.aitken_has_previous is not None: - mapping.aitken_has_previous.zero_() + wp.launch( + _zero_global_proxy_values_masked_kernel, + dim=mapping.proxy_ids_global.shape[0], + inputs=[mapping.proxy_ids_global, entity_world, world_mask, mapping.coupling_forces], + device=self.model.device, + ) + for values in (mapping.coupling_forces_previous, mapping.aitken_residual_previous): + if values is not None: + wp.launch( + _zero_proxy_row_values_masked_kernel, + dim=mapping.proxy_ids_global.shape[0], + inputs=[mapping.proxy_ids_global, entity_world, world_mask, values], + device=self.model.device, + ) if mapping.proxy_qd_before is not None: - mapping.proxy_qd_before.zero_() + wp.launch( + _zero_local_proxy_values_masked_kernel, + dim=mapping.proxy_ids_global.shape[0], + inputs=[ + mapping.proxy_ids_global, + mapping.proxy_ids_local, + entity_world, + world_mask, + mapping.proxy_qd_before, + ], + device=self.model.device, + ) for config in self._proxy_collision_configs.values(): config.collide_counter = 0 if config.contacts is not None: diff --git a/newton/_src/solvers/featherstone/kernels.py b/newton/_src/solvers/featherstone/kernels.py index b02c566f89..7e03f11cf2 100644 --- a/newton/_src/solvers/featherstone/kernels.py +++ b/newton/_src/solvers/featherstone/kernels.py @@ -415,7 +415,9 @@ def jcalc_tau( # w = joint_qd[dof_start + i] # r = joint_q[coord_start + i] - tau[dof_start + i] = -wp.dot(S_s, body_f_s) + joint_f[dof_start + i] + j = dof_start + i + passive_f = -joint_damping[j] * joint_qd[j] + tau[j] = -wp.dot(S_s, body_f_s) + joint_f[j] + passive_f # tau -= w * target_kd - r * target_ke return diff --git a/newton/_src/solvers/flags.py b/newton/_src/solvers/flags.py deleted file mode 100644 index 71f05a4145..0000000000 --- a/newton/_src/solvers/flags.py +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers -# SPDX-License-Identifier: Apache-2.0 - -"""Solver flags.""" - -import warnings -from enum import EnumMeta, IntEnum - -from ..sim.enums import ModelFlags - - -class _DeprecatedSolverNotifyFlagsMeta(EnumMeta): - def __getattribute__(cls, name: str): - value = super().__getattribute__(name) - if not name.startswith("_"): - member_map = super().__getattribute__("_member_map_") - if name in member_map: - _warn_solver_notify_flags_deprecated() - return value - - def __call__(cls, *args, **kwargs): - _warn_solver_notify_flags_deprecated() - return super().__call__(*args, **kwargs) - - -def _warn_solver_notify_flags_deprecated() -> None: - warnings.warn( - "SolverNotifyFlags is deprecated, use ModelFlags instead.", - DeprecationWarning, - stacklevel=3, - ) - - -class SolverNotifyFlags(IntEnum, metaclass=_DeprecatedSolverNotifyFlagsMeta): - """Deprecated alias for :class:`~newton.ModelFlags`. - - .. deprecated:: 1.3 - Use :class:`~newton.ModelFlags` instead. - """ - - JOINT_PROPERTIES = ModelFlags.JOINT_PROPERTIES.value - JOINT_DOF_PROPERTIES = ModelFlags.JOINT_DOF_PROPERTIES.value - BODY_PROPERTIES = ModelFlags.BODY_PROPERTIES.value - BODY_INERTIAL_PROPERTIES = ModelFlags.BODY_INERTIAL_PROPERTIES.value - SHAPE_PROPERTIES = ModelFlags.SHAPE_PROPERTIES.value - MODEL_PROPERTIES = ModelFlags.MODEL_PROPERTIES.value - CONSTRAINT_PROPERTIES = ModelFlags.CONSTRAINT_PROPERTIES.value - TENDON_PROPERTIES = ModelFlags.TENDON_PROPERTIES.value - ACTUATOR_PROPERTIES = ModelFlags.ACTUATOR_PROPERTIES.value - ALL = ModelFlags.ALL.value - - -__all__ = [ - "SolverNotifyFlags", -] diff --git a/newton/_src/solvers/implicit_mpm/implicit_mpm_model.py b/newton/_src/solvers/implicit_mpm/implicit_mpm_model.py index 19b3234d22..b82bf39415 100644 --- a/newton/_src/solvers/implicit_mpm/implicit_mpm_model.py +++ b/newton/_src/solvers/implicit_mpm/implicit_mpm_model.py @@ -304,6 +304,9 @@ def __init__(self, model: newton.Model, options: SolverImplicitMPM.Config): self.air_drag = float(options.air_drag) """Drag for the background air""" + self._separate_worlds = bool(options.separate_worlds and model.world_count > 1) + """Whether colliders and particles are isolated by Newton world.""" + self.collider = Collider() """Collider struct""" @@ -403,14 +406,15 @@ def _refresh_particle_flags_and_extrema(self): self.has_viscosity = False self.has_dilatancy = False - def notify_collider_changed(self): + def notify_collider_changed(self, body_mass: np.ndarray | None = None): """Refresh cached extrema for collider parameters. Tracks the minimum collider mass to determine whether compliant colliders are present and to enable/disable related computations. """ body_ids = self.collider.collider_body_index.numpy() - body_mass = self.collider_body_mass.numpy() + if body_mass is None: + body_mass = self.collider_body_mass.numpy() dynamic_body_ids = body_ids[body_ids >= 0] dynamic_body_ids = dynamic_body_ids[body_mass[dynamic_body_ids] > 0.0] dynamic_body_masses = body_mass[dynamic_body_ids] @@ -439,6 +443,7 @@ def setup_collider( body_mass: wp.array | None = None, body_inv_inertia: wp.array | None = None, body_q: wp.array | None = None, + collider_world_ids: list[int] | None = None, ): """Initialize collider parameters and defaults from inputs. @@ -446,8 +451,9 @@ def setup_collider( properties (thickness, friction, adhesion, projection threshold). By default, this will setup collisions against all collision shapes in the model with flag `newton.ShapeFlag.COLLIDE_PARTICLES`. - Rigid body colliders will be treated as kinematic if their mass is zero; for all model bodies to be treated as kinematic, - pass ``body_mass=wp.zeros_like(model.body_mass)``. + Rigid body colliders will be treated as kinematic if their effective mass is zero. When ``body_mass`` is omitted, + bodies flagged with :attr:`newton.BodyFlags.KINEMATIC` have zero effective mass regardless of their stored mass. + An explicit ``body_mass`` array is authoritative. For any collider index `i`, only one of ``collider_meshes[i]`` and ``collider_body_ids`` may not be `None`. If material properties are not provided for a collider, but a body index is provided, @@ -461,61 +467,219 @@ def setup_collider( collider_adhesion: Per-mesh adhesion (Pa). collider_projection_threshold: Per-mesh projection threshold, i.e. how far below the surface the particle may be before it is projected out. (m) - collider_particle_ids: For deformable mesh colliders, model particle ids corresponding to each mesh vertex. + collider_particle_ids: For deformable mesh colliders, solver-model particle IDs corresponding to each + mesh vertex. These IDs cannot be combined with an external ``model``. In isolated multi-world mode, + every ID must belong to the collider's local world; global deformable colliders are rejected. model: The model to read collider properties from. Default to self.model. body_com: For dynamic colliders, per-body center of mass. Default to model.body_com. - body_mass: For dynamic colliders, per-body mass. Default to model.body_mass. + body_mass: For dynamic colliders, per-body effective mass. By default, use ``model.body_mass`` with + zero mass for bodies flagged with :attr:`newton.BodyFlags.KINEMATIC`. body_inv_inertia: For dynamic colliders, per-body inverse inertia. Default to model.body_inv_inertia. body_q: For dynamic colliders, per-body initial transform. Default to model.body_q. + collider_world_ids: Per-collider Newton world IDs. Custom meshes default to global + (``-1``), while body-backed colliders infer their body's world. + + Raises: + ValueError: If collider inputs are inconsistent, world IDs are invalid, an isolated external model has a + different world count, a global body-backed collider is dynamic, or deformable collider particle + ownership cannot be mapped safely to the solver model and world. """ if model is None: model = self.model + elif self._separate_worlds and model is not self.model and model.world_count != self.model.world_count: + raise ValueError( + "An external collider model must have the same world_count as the isolated solver model; " + f"got {model.world_count} and {self.model.world_count}." + ) - if collider_body_ids is None: - if collider_meshes is None: - collider_body_ids = [ - body_id - for body_id in range(-1, model.body_count) - if len(_get_body_collision_shapes(model, body_id)) > 0 - ] + collider_meshes = None if collider_meshes is None else list(collider_meshes) + collider_body_ids = None if collider_body_ids is None else list(collider_body_ids) + collider_thicknesses = None if collider_thicknesses is None else list(collider_thicknesses) + collider_friction = None if collider_friction is None else list(collider_friction) + collider_adhesion = None if collider_adhesion is None else list(collider_adhesion) + collider_projection_threshold = ( + None if collider_projection_threshold is None else list(collider_projection_threshold) + ) + collider_particle_ids = None if collider_particle_ids is None else list(collider_particle_ids) + supplied_world_ids = None if collider_world_ids is None else list(collider_world_ids) + + shape_world = None + body_world = None + + def get_shape_world(): + nonlocal shape_world + if shape_world is None: + shape_world = ( + model.shape_world.numpy() + if model.shape_world is not None + else np.full(model.shape_count, -1, dtype=int) + ) + return shape_world + + def get_body_world(): + nonlocal body_world + if body_world is None: + body_world = ( + model.body_world.numpy() + if model.body_world is not None + else np.full(model.body_count, -1, dtype=int) + ) + return body_world + + default_discovery = collider_meshes is None and collider_body_ids is None + collider_shapes = [] + inferred_world_ids = [] + if default_discovery: + collider_meshes = [] + collider_body_ids = [] + + static_shapes = _get_body_collision_shapes(model, -1) + if self._separate_worlds: + static_collider_worlds = sorted({int(world) for world in get_shape_world()[static_shapes]}) else: - collider_body_ids = [None] * len(collider_meshes) - if collider_meshes is None: - collider_meshes = [None] * len(collider_body_ids) + static_collider_worlds = [-1] if len(static_shapes) > 0 else [] + for world_id in static_collider_worlds: + collider_meshes.append(None) + collider_body_ids.append(-1) + collider_shapes.append( + static_shapes[get_shape_world()[static_shapes] == world_id] + if self._separate_worlds + else static_shapes + ) + inferred_world_ids.append(world_id) - for collider_id, (mesh, body_id) in enumerate(zip(collider_meshes, collider_body_ids, strict=True)): - if mesh is None: - if body_id is None: - raise ValueError( - f"Either a mesh or a body_id must be provided for each collider; collider {collider_id} is missing both" - ) - elif body_id is not None: + for body_id in range(model.body_count): + shapes = _get_body_collision_shapes(model, body_id) + if len(shapes) == 0: + continue + collider_meshes.append(None) + collider_body_ids.append(body_id) + collider_shapes.append(shapes) + inferred_world_ids.append(int(get_body_world()[body_id]) if self._separate_worlds else -1) + else: + if collider_body_ids is None: + collider_body_ids = [None] * len(collider_meshes) + elif collider_meshes is None: + collider_meshes = [None] * len(collider_body_ids) + elif len(collider_meshes) != len(collider_body_ids): raise ValueError( - f"Either a mesh or a body_id must be provided for each collider; collider {collider_id} provides both" + "collider_meshes and collider_body_ids must have the same length; " + f"got {len(collider_meshes)} and {len(collider_body_ids)}." ) + collider_shapes = [None] * len(collider_body_ids) collider_count = len(collider_body_ids) - if collider_thicknesses is None: - collider_thicknesses = [None] * collider_count - if collider_projection_threshold is None: - collider_projection_threshold = [None] * collider_count - if collider_friction is None: - collider_friction = [None] * collider_count - if collider_adhesion is None: - collider_adhesion = [None] * collider_count - if collider_particle_ids is None: - collider_particle_ids = [None] * collider_count - - assert len(collider_body_ids) == len(collider_thicknesses) - assert len(collider_body_ids) == len(collider_projection_threshold) - assert len(collider_body_ids) == len(collider_friction) - assert len(collider_body_ids) == len(collider_adhesion) - assert len(collider_body_ids) == len(collider_particle_ids) + def require_aligned(name, values, default=None): + if values is None: + return [default] * collider_count + if len(values) != collider_count: + raise ValueError(f"{name} must have one value per collider ({collider_count}); got {len(values)}.") + return values + + collider_meshes = require_aligned("collider_meshes", collider_meshes) + collider_thicknesses = require_aligned("collider_thicknesses", collider_thicknesses) + collider_projection_threshold = require_aligned("collider_projection_threshold", collider_projection_threshold) + collider_friction = require_aligned("collider_friction", collider_friction) + collider_adhesion = require_aligned("collider_adhesion", collider_adhesion) + collider_particle_ids = require_aligned("collider_particle_ids", collider_particle_ids) + supplied_world_ids = require_aligned("collider_world_ids", supplied_world_ids) + + def validate_world_id(world_id, collider_id): + if not isinstance(world_id, (int, np.integer)): + raise ValueError(f"Invalid collider world ID {world_id!r} for collider {collider_id}.") + world_id = int(world_id) + if world_id < -1 or world_id >= self.model.world_count: + raise ValueError( + f"Invalid collider world ID {world_id} for collider {collider_id}; expected -1 or an ID in " + f"[0, {self.model.world_count})." + ) + return world_id + + if default_discovery: + collider_world_ids = [] + for collider_id, raw_inferred_world_id in enumerate(inferred_world_ids): + inferred_world_id = validate_world_id(raw_inferred_world_id, collider_id) + supplied_world_id = supplied_world_ids[collider_id] + if supplied_world_id is not None: + supplied_world_id = validate_world_id(supplied_world_id, collider_id) + if supplied_world_id != inferred_world_id: + raise ValueError( + f"Collider world ID {supplied_world_id} for collider {collider_id} does not match " + f"its inferred world ID {inferred_world_id}." + ) + collider_world_ids.append(inferred_world_id) + else: + collider_world_ids = [] + static_shapes = None + for collider_id, (mesh, raw_body_id, requested_world_id) in enumerate( + zip(collider_meshes, collider_body_ids, supplied_world_ids, strict=True) + ): + if mesh is None and raw_body_id is None: + raise ValueError( + f"Either a mesh or a body_id must be provided for each collider; collider {collider_id} is missing both" + ) + if mesh is not None and raw_body_id is not None: + raise ValueError( + f"Either a mesh or a body_id must be provided for each collider; collider {collider_id} provides both" + ) + + if raw_body_id is None: + world_id = -1 if requested_world_id is None else requested_world_id + else: + try: + body_id = int(raw_body_id) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Invalid collider body ID {raw_body_id!r} for collider {collider_id}." + ) from exc + if body_id < -1 or body_id >= model.body_count: + raise ValueError( + f"Invalid collider body ID {body_id} for collider {collider_id}; expected -1 or an ID " + f"in [0, {model.body_count})." + ) + collider_body_ids[collider_id] = body_id + + if body_id >= 0: + if self._separate_worlds: + world_id = validate_world_id(get_body_world()[body_id], collider_id) + else: + world_id = -1 + if requested_world_id is not None and self._separate_worlds: + validated_world_id = validate_world_id(requested_world_id, collider_id) + if validated_world_id != world_id: + raise ValueError( + f"Collider world ID {validated_world_id} for collider {collider_id} does not " + f"match body {body_id}'s inferred world ID {world_id}." + ) + elif requested_world_id is not None: + world_id = requested_world_id + collider_shapes[collider_id] = _get_body_collision_shapes(model, body_id) + else: + if static_shapes is None: + static_shapes = _get_body_collision_shapes(model, -1) + if self._separate_worlds: + if requested_world_id is None: + static_world_ids = sorted({int(world) for world in get_shape_world()[static_shapes]}) + if len(static_world_ids) != 1: + raise ValueError( + "An explicit body -1 collider is ambiguous in isolated multi-world mode; " + "supply its collider world ID or use default discovery/custom meshes." + ) + world_id = static_world_ids[0] + else: + world_id = validate_world_id(requested_world_id, collider_id) + collider_shapes[collider_id] = static_shapes[get_shape_world()[static_shapes] == world_id] + else: + world_id = -1 if requested_world_id is None else requested_world_id + collider_shapes[collider_id] = static_shapes + + collider_world_ids.append(validate_world_id(world_id, collider_id)) if body_com is None: body_com = model.body_com + body_mass_is_explicit = body_mass is not None if body_mass is None: body_mass = model.body_mass if body_inv_inertia is None: @@ -523,17 +687,30 @@ def setup_collider( if body_q is None: body_q = model.body_q + effective_body_mass = body_mass.numpy().copy() + if not body_mass_is_explicit and model.body_flags is not None: + body_flags = model.body_flags.numpy() + kinematic_bodies = (body_flags & int(newton.BodyFlags.KINEMATIC)) != 0 + if np.any(kinematic_bodies & (effective_body_mass != 0.0)): + effective_body_mass[kinematic_bodies] = 0.0 + body_mass = wp.array(effective_body_mass, dtype=body_mass.dtype, device=body_mass.device) + if self._separate_worlds: + for collider_id, (body_id, world_id) in enumerate(zip(collider_body_ids, collider_world_ids, strict=True)): + if world_id == -1 and body_id is not None and body_id >= 0 and effective_body_mass[body_id] > 0.0: + raise ValueError( + f"Collider {collider_id} is a global dynamic collider backed by body {body_id}, which would " + "couple isolated worlds. Replicate it into each world, make it static or kinematic, or disable " + "Config.separate_worlds." + ) + # count materials and shapes material_count = 1 # default material - body_shapes = {} collider_material_ids = [] - for body_id in collider_body_ids: + for body_id, shapes in zip(collider_body_ids, collider_shapes, strict=True): if body_id is not None: - shapes = _get_body_collision_shapes(model, body_id) if len(shapes) == 0: - raise ValueError(f"Body {body_id} has no collision shapes") + raise ValueError(f"Body {body_id} has no collision shapes for its collider world ID") - body_shapes[body_id] = shapes collider_material_ids.append(list(range(material_count, material_count + len(shapes)))) material_count += len(shapes) else: @@ -575,7 +752,7 @@ def assign_collider_material(material_id: int, collider_id: int): if body_id is not None: for material_id, shape_margin, shape_friction in zip( collider_material_ids[collider_id], - *_get_shape_collision_materials(model, body_shapes[body_id]), + *_get_shape_collision_materials(model, collider_shapes[collider_id]), strict=True, ): # use material from shapes as default @@ -590,6 +767,16 @@ def assign_collider_material(material_id: int, collider_id: int): max((material_thickness[material_id] for material_id in collider_material_ids[collider_id]), default=0.0) for collider_id in range(collider_count) ] + has_deformable_colliders = any(particle_ids is not None for particle_ids in collider_particle_ids) + if has_deformable_colliders and model is not self.model: + raise ValueError( + "collider_particle_ids are solver-state particle indices and may only be used with the solver model; " + "external collider models are not supported" + ) + + solver_particle_world = ( + self.model.particle_world.numpy() if self._separate_worlds and has_deformable_colliders else None + ) collider_particle_offsets = [0] collider_particle_id_chunks = [] for collider_id, particle_ids in enumerate(collider_particle_ids): @@ -614,9 +801,24 @@ def assign_collider_material(material_id: int, collider_id: int): f"but collider mesh has {vertex_count} vertices" ) if particle_ids_np.size and ( - np.min(particle_ids_np) < 0 or np.max(particle_ids_np) >= model.particle_count + np.min(particle_ids_np) < 0 or np.max(particle_ids_np) >= self.model.particle_count ): - raise ValueError(f"collider_particle_ids[{collider_id}] contains particle ids outside the model") + raise ValueError(f"collider_particle_ids[{collider_id}] contains particle ids outside the solver model") + + if self._separate_worlds: + collider_world_id = collider_world_ids[collider_id] + if collider_world_id < 0: + raise ValueError( + f"collider_particle_ids[{collider_id}] cannot define a global deformable collider across " + "isolated worlds; assign the collider to one local world" + ) + + particle_world_ids = np.unique(solver_particle_world[particle_ids_np]) + if np.any(particle_world_ids != collider_world_id): + raise ValueError( + f"collider_particle_ids[{collider_id}] must reference particles in collider world " + f"{collider_world_id}; found particle world IDs {particle_world_ids.tolist()}" + ) collider_particle_id_chunks.append(particle_ids_np) collider_particle_offsets.append(collider_particle_offsets[-1] + particle_ids_np.shape[0]) @@ -631,7 +833,10 @@ def assign_collider_material(material_id: int, collider_id: int): # Create device arrays with wp.ScopedDevice(self.model.device): # Create collider meshes from bodies if necessary - face_material_ids = [[]] + packed_body_ids = [] + face_material_ids = [] + collider_face_offsets = [] + face_offset = 0 for collider_id in range(collider_count): body_index = collider_body_ids[collider_id] @@ -640,24 +845,47 @@ def assign_collider_material(material_id: int, collider_id: int): # This may not correspond to the model's body -1, but as far as the collision kernels # are concerned, it does not matter. - collider_body_ids[collider_id] = -1 + packed_body_ids.append(-1) material_id = collider_material_ids[collider_id][0] face_count = collider_meshes[collider_id].indices.shape[0] // 3 mesh_face_material_ids = np.full(face_count, material_id, dtype=int) else: collider_meshes[collider_id], mesh_face_material_ids = _create_body_collider_mesh( - model, body_shapes[body_index], collider_material_ids[collider_id] + model, collider_shapes[collider_id], collider_material_ids[collider_id] ) + packed_body_ids.append(body_index) + face_count = collider_meshes[collider_id].indices.shape[0] // 3 face_material_ids.append(mesh_face_material_ids) + collider_face_offsets.append(face_offset) + face_offset += face_count + + global_collider_ids = [ + collider_id for collider_id, collider_world_id in enumerate(collider_world_ids) if collider_world_id < 0 + ] + world_collider_ids = [] + world_collider_offsets = [0] + for world_id in range(self.model.world_count): + world_collider_ids.extend(global_collider_ids) + world_collider_ids.extend( + collider_id + for collider_id, collider_world_id in enumerate(collider_world_ids) + if collider_world_id == world_id + ) + world_collider_offsets.append(len(world_collider_ids)) - self.collider.collider_body_index = wp.array(collider_body_ids, dtype=int) + self.collider.collider_body_index = wp.array(packed_body_ids, dtype=int) self.collider.collider_particle_offsets = wp.array(collider_particle_offsets, dtype=int) self.collider.collider_particle_ids = wp.array(flat_collider_particle_ids, dtype=int) self.collider.collider_mesh = wp.array([collider.id for collider in collider_meshes], dtype=wp.uint64) self.collider.collider_max_thickness = wp.array(collider_max_thickness, dtype=float) + self.collider.collider_world = wp.array(collider_world_ids, dtype=int) + self.collider.collider_face_offset = wp.array(collider_face_offsets, dtype=int) + self.collider.world_collider_ids = wp.array(world_collider_ids, dtype=int) + self.collider.world_collider_offsets = wp.array(world_collider_offsets, dtype=int) - self.collider.face_material_index = wp.array(np.concatenate(face_material_ids), dtype=int) + all_face_material_ids = np.concatenate(face_material_ids) if face_material_ids else np.empty(0, dtype=int) + self.collider.face_material_index = wp.array(all_face_material_ids, dtype=int) self.collider.material_thickness = wp.array(material_thickness, dtype=float) self.collider.material_friction = wp.array(material_friction, dtype=float) @@ -676,7 +904,7 @@ def assign_collider_material(material_id: int, collider_id: int): ] self._refresh_particle_flags_and_extrema() - self.notify_collider_changed() + self.notify_collider_changed(effective_body_mass) @property def has_compliant_particles(self): diff --git a/newton/_src/solvers/implicit_mpm/implicit_mpm_solver_kernels.py b/newton/_src/solvers/implicit_mpm/implicit_mpm_solver_kernels.py index 01e5a24b1d..3457ab4fb5 100644 --- a/newton/_src/solvers/implicit_mpm/implicit_mpm_solver_kernels.py +++ b/newton/_src/solvers/implicit_mpm/implicit_mpm_solver_kernels.py @@ -1,8 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers # SPDX-License-Identifier: Apache-2.0 +import math from typing import Any +import numpy as np import warp as wp import warp.fem as fem import warp.sparse as wps @@ -881,6 +883,120 @@ def clamp_coordinates( coords[i] = wp.min(wp.max(coords[i], wp.vec3(0.0)), wp.vec3(1.0)) +@wp.kernel +def build_active_particle_mask( + particle_q: wp.array[wp.vec3], + particle_flags: wp.array[wp.int32], + point_mask: wp.array[wp.int32], +): + particle_index = wp.tid() + position = particle_q[particle_index] + position_is_finite = wp.isfinite(position[0]) and wp.isfinite(position[1]) and wp.isfinite(position[2]) + point_mask[particle_index] = wp.where( + (particle_flags[particle_index] & newton.ParticleFlags.ACTIVE) != 0 and position_is_finite, + wp.int32(1), + wp.int32(0), + ) + + +@wp.kernel +def record_volume_rebuild_status(status: wp.array[wp.uint32], accumulated_status: wp.array[wp.uint32]): + """Retain and report rebuild failures without a host synchronization.""" + rebuild_status = status[0] + new_status = rebuild_status & ~accumulated_status[0] + if new_status != wp.uint32(0): + accumulated_status[0] = accumulated_status[0] | rebuild_status + wp.printf("Warning: Implicit MPM sparse grid rebuild failed with status %u.\n", rebuild_status) + + +@wp.func +def reset_mpm_world_is_selected(world: int, world_mask: wp.array[wp.bool]): + selected = bool(False) + global_world_index = world_mask.shape[0] - 1 + if world >= 0 and world < global_world_index: + selected = world_mask[world] + elif world == -1: + selected = world_mask[global_world_index] + return selected + + +@wp.kernel +def reset_mpm_particle_history( + particle_world: wp.array[wp.int32], + world_mask: wp.array[wp.bool], + particle_elastic_strain: wp.array[wp.mat33], + particle_transform: wp.array[wp.mat33], + particle_qd_grad: wp.array[wp.mat33], + particle_stress: wp.array[wp.mat33], + particle_Jp: wp.array[float], +): + """Reset implicit MPM history for selected local or shared particles.""" + particle_index = wp.tid() + world = particle_world[particle_index] + if reset_mpm_world_is_selected(world, world_mask): + identity = wp.identity(n=3, dtype=float) + particle_elastic_strain[particle_index] = identity + particle_transform[particle_index] = identity + particle_qd_grad[particle_index] = wp.mat33(0.0) + particle_stress[particle_index] = wp.mat33(0.0) + particle_Jp[particle_index] = 1.0 + + +@wp.kernel +def reset_mpm_collider_history( + body_world: wp.array[wp.int32], + world_mask: wp.array[wp.bool], + body_q: wp.array[wp.transform], + body_q_prev: wp.array[wp.transform], +): + """Refresh previous collider poses for selected local or shared bodies.""" + body_index = wp.tid() + world = body_world[body_index] + if reset_mpm_world_is_selected(world, world_mask): + body_q_prev[body_index] = body_q[body_index] + + +@wp.kernel(module="unique") +def reset_mpm_point_warmstart( + particle_world: wp.array[wp.int32], + world_mask: wp.array[wp.bool], + values: wp.array[Any], +): + """Clear particle-backed warm starts for selected local or shared particles.""" + particle_index = wp.tid() + world = particle_world[particle_index] + if reset_mpm_world_is_selected(world, world_mask): + values[particle_index] = values.dtype(0.0) + + +wp.overload(reset_mpm_point_warmstart, {"values": wp.array[wp.vec3]}) +wp.overload(reset_mpm_point_warmstart, {"values": wp.array[vec6]}) + + +@wp.kernel(module="unique") +def reset_mpm_grid_warmstart( + world_mask: wp.array[wp.bool], + environment_offsets: wp.array[int], + environment_node_indices: wp.array[int], + values: wp.array[Any], +): + """Clear whole-space warm-start values for selected environments.""" + partition_index = wp.tid() + environment_count = world_mask.shape[0] - 1 + if partition_index >= environment_offsets[environment_count]: + return + + environment = wp.lower_bound(environment_offsets, partition_index + 1) - 1 + if environment >= 0 and environment < environment_count and world_mask[environment]: + space_node_index = environment_node_indices[partition_index] + if space_node_index >= 0 and space_node_index < values.shape[0]: + values[space_node_index] = values.dtype(0.0) + + +wp.overload(reset_mpm_grid_warmstart, {"values": wp.array[wp.vec3]}) +wp.overload(reset_mpm_grid_warmstart, {"values": wp.array[vec6]}) + + @wp.kernel def pad_voxels(particle_q: wp.array[wp.vec3i], padded_q: wp.array4d[wp.vec3i]): pid = wp.tid() @@ -896,7 +1012,120 @@ def positive_modn(x: int, n: int): return (x % n + n) % n -def allocate_by_voxels(particle_q, voxel_size, padding_voxels: int = 0): +def _rebuild_capacity( + particle_q, + voxel_size, + ratio: float, + max_active_voxels: int, + point_mask=None, + *, + point_environment=None, + environment_count: int | None = None, + guard_cells: int = 3, + temporary_store=None, + max_leaf_node_count: int = -1, + max_lower_node_count: int = -1, + max_upper_node_count: int = -1, +) -> dict[str, int]: + """Estimate rebuildable-volume capacities (active voxels + NanoVDB node counts). + + Automatic leaf-node storage reserves ``max_active_voxels`` entries. The + lower/upper internal nodes each span 16x and 32x more cells, so their counts + are typically small; automatic capacities are estimated from one throwaway + build of the current particles scaled by ``ratio`` for spreading headroom. + Explicit node capacities allow applications with known spatial bounds to + budget each NanoVDB hierarchy level independently. Rebuild status reports + when any reserved capacity is exceeded. + """ + if point_environment is None: + initial = wp.Volume.allocate_by_voxels( + voxel_points=particle_q, + voxel_size=voxel_size, + point_mask=point_mask, + ) + ijk = initial.get_voxels().numpy() + else: + if environment_count is None: + raise ValueError("environment_count is required with point_environment") + initial = fem.Nanogrid.from_environment_voxels( + particle_q, + point_environment, + environment_count, + point_mask=point_mask, + guard_cells=guard_cells, + voxel_size=voxel_size, + temporary_store=temporary_store, + device=particle_q.device, + ) + ijk = initial.cell_grid.get_voxels().numpy() + + if ijk.shape[0] == 0: + automatic_lower = min(max_active_voxels, 8) + automatic_upper = min(max_active_voxels, 4) + else: + lower = np.unique(np.floor_divide(ijk, 8 * 16), axis=0).shape[0] + upper = np.unique(np.floor_divide(ijk, 8 * 16 * 32), axis=0).shape[0] + automatic_lower = min(max_active_voxels, max(8, math.ceil(lower * ratio))) + automatic_upper = min(max_active_voxels, max(4, math.ceil(upper * ratio))) + + leaf_capacity = max_active_voxels if max_leaf_node_count == -1 else max_leaf_node_count + lower_capacity = min(automatic_lower, leaf_capacity) if max_lower_node_count == -1 else max_lower_node_count + upper_capacity = min(automatic_upper, lower_capacity) if max_upper_node_count == -1 else max_upper_node_count + + if not upper_capacity <= lower_capacity <= leaf_capacity <= max_active_voxels: + raise ValueError( + "Implicit MPM sparse-grid capacity hierarchy must satisfy " + "max_upper_node_count <= max_lower_node_count <= max_leaf_node_count " + "<= max_active_cell_count after resolving automatic values; got " + f"{upper_capacity} <= {lower_capacity} <= {leaf_capacity} <= {max_active_voxels}." + ) + + return { + "max_active_voxels": max_active_voxels, + "max_leaf_nodes": leaf_capacity, + "max_lower_nodes": lower_capacity, + "max_upper_nodes": upper_capacity, + } + + +def allocate_by_voxels( + particle_q, + voxel_size, + padding_voxels: int = 0, + rebuildable: bool = False, + max_active_voxels: int | None = None, + capacity_ratio: float = 16.0, + status=None, + point_mask=None, + max_leaf_node_count: int = -1, + max_lower_node_count: int = -1, + max_upper_node_count: int = -1, +): + if rebuildable: + # Persistent capacity-sized volume refreshed in place each step so the sparse + # grid build is CUDA-graph-capturable. Padding is unsupported here. + capacity = max_active_voxels if max_active_voxels and max_active_voxels > 0 else particle_q.shape[0] + kwargs = _rebuild_capacity( + particle_q, + voxel_size, + capacity_ratio, + capacity, + point_mask=point_mask, + max_leaf_node_count=max_leaf_node_count, + max_lower_node_count=max_lower_node_count, + max_upper_node_count=max_upper_node_count, + ) + if status is not None: + kwargs["status"] = status + if point_mask is not None: + kwargs["point_mask"] = point_mask + return wp.Volume.allocate_by_voxels( + voxel_points=particle_q, + voxel_size=voxel_size, + rebuildable=True, + **kwargs, + ) + volume = wp.Volume.allocate_by_voxels( voxel_points=particle_q.flatten(), voxel_size=voxel_size, @@ -917,6 +1146,16 @@ def allocate_by_voxels(particle_q, voxel_size, padding_voxels: int = 0): return volume +def voxel_coordinates(particle_q: wp.array[wp.vec3], voxel_size: float, padding_voxels: int = 0) -> wp.array[wp.vec3i]: + if particle_q.shape[0] == 0: + return wp.empty(0, dtype=wp.vec3i, device=particle_q.device) + + volume = allocate_by_voxels(particle_q, voxel_size, padding_voxels=padding_voxels) + voxels = wp.empty(volume.get_voxel_count(), dtype=wp.vec3i, device=particle_q.device) + volume.get_voxels(voxels) + return voxels + + @wp.kernel def node_color( space_node_indices: wp.array[int], @@ -1053,6 +1292,25 @@ def mark_active_cells( active_cells[s_grid.element_index] = 1 +@fem.integrand +def mark_active_cells_by_environment( + s: fem.Sample, + domain: fem.Domain, + positions: wp.array[wp.vec3], + particle_flags: wp.array[int], + particle_environment: wp.array[int], + active_cells: wp.array[int], +): + if ~particle_flags[s.qp_index] & newton.ParticleFlags.ACTIVE: + return + + x = positions[s.qp_index] + s_grid = fem.lookup(domain, x, int(particle_environment[s.qp_index])) + + if s_grid.element_index != fem.NULL_ELEMENT_INDEX: + active_cells[s_grid.element_index] = 1 + + @wp.kernel(module="unique") def scatter_field_dof_values( space_node_indices: wp.array[int], diff --git a/newton/_src/solvers/implicit_mpm/rasterized_collisions.py b/newton/_src/solvers/implicit_mpm/rasterized_collisions.py index 3ad99c74df..bdf0d5b8b6 100644 --- a/newton/_src/solvers/implicit_mpm/rasterized_collisions.py +++ b/newton/_src/solvers/implicit_mpm/rasterized_collisions.py @@ -37,6 +37,9 @@ _NULL_COLLIDER_ID = -1 """Indicator for no collider""" +_ALL_COLLIDER_WORLDS = -2 +"""Environment sentinel that queries every collider in stable order.""" + @wp.struct class Collider: @@ -57,6 +60,18 @@ class Collider: collider_particle_ids: wp.array[int] """Model particle index for each deformable collider mesh vertex. Shape (sum(deformable mesh vertex counts),)""" + collider_world: wp.array[int] + """Newton world ID of each stable collider. Shape (collider_count,).""" + + collider_face_offset: wp.array[int] + """Start of each stable collider's faces in ``face_material_index``. Shape (collider_count,).""" + + world_collider_ids: wp.array[int] + """Stable collider IDs affecting each world, grouped by world.""" + + world_collider_offsets: wp.array[int] + """Offsets into ``world_collider_ids`` for each world. Shape (world_count + 1,).""" + face_material_index: wp.array[int] """Material index for each collider mesh face. Shape (sum(mesh.face_count for mesh in meshes),)""" @@ -159,72 +174,106 @@ def get_average_face_normal( @wp.func -def collision_sdf( +def _query_collider_sdf( x: wp.vec3, collider: Collider, body_q: wp.array[wp.transform], - body_qd: wp.array[wp.spatial_vector], - body_q_prev: wp.array[wp.transform], - dt: float, + stable_collider_id: int, ): - min_sdf = float(_INFINITY) + mesh = collider.collider_mesh[stable_collider_id] + thickness = collider.collider_max_thickness[stable_collider_id] + body_id = collider.collider_body_index[stable_collider_id] + + if body_id >= 0: + b_pos = wp.transform_get_translation(body_q[body_id]) + b_rot = wp.transform_get_rotation(body_q[body_id]) + x_local = wp.quat_rotate_inv(b_rot, x - b_pos) + else: + x_local = x + + max_dist = collider.query_max_dist + thickness + + if wp.static(_SDF_SIGN_FROM_AVERAGE_NORMAL): + query = wp.mesh_query_point_no_sign(mesh, x_local, max_dist) + else: + query = wp.mesh_query_point(mesh, x_local, max_dist) + + query_result = query.result + sdf = float(_INFINITY) sdf_grad = wp.vec3(0.0) sdf_vel = wp.vec3(0.0) closest_point = wp.vec3(0.0) - collider_id = int(_NULL_COLLIDER_ID) material_id = int(0) # default material, always valid - # Find closest collider - global_face_id = int(0) - for m in range(collider.collider_mesh.shape[0]): - mesh = collider.collider_mesh[m] - thickness = collider.collider_max_thickness[m] - body_id = collider.collider_body_index[m] + if query_result: + cp = wp.mesh_eval_position(mesh, query.face, query.u, query.v) - if body_id >= 0: - b_pos = wp.transform_get_translation(body_q[body_id]) - b_rot = wp.transform_get_rotation(body_q[body_id]) - x_local = wp.quat_rotate_inv(b_rot, x - b_pos) + if wp.static(_SDF_SIGN_FROM_AVERAGE_NORMAL): + face_normal = get_average_face_normal(mesh, cp) + sign = wp.where(wp.dot(face_normal, x_local - cp) > 0.0, 1.0, -1.0) else: - x_local = x + face_normal = wp.mesh_eval_face_normal(mesh, query.face) + sign = query.sign - max_dist = collider.query_max_dist + thickness + mesh_material_id = collider.face_material_index[collider.collider_face_offset[stable_collider_id] + query.face] + thickness = collider.material_thickness[mesh_material_id] - if wp.static(_SDF_SIGN_FROM_AVERAGE_NORMAL): - query = wp.mesh_query_point_no_sign(mesh, x_local, max_dist) + offset = x_local - cp + d = wp.length(offset) * sign + sdf = d - thickness + + if wp.abs(d) < _CLOSEST_POINT_NORMAL_EPSILON: + sdf_grad = face_normal else: - query = wp.mesh_query_point(mesh, x_local, max_dist) + sdf_grad = wp.normalize(offset) * sign - if query.result: - cp = wp.mesh_eval_position(mesh, query.face, query.u, query.v) + sdf_vel = wp.mesh_eval_velocity(mesh, query.face, query.u, query.v) + closest_point = cp + material_id = mesh_material_id - if wp.static(_SDF_SIGN_FROM_AVERAGE_NORMAL): - face_normal = get_average_face_normal(mesh, cp) - sign = wp.where(wp.dot(face_normal, x_local - cp) > 0.0, 1.0, -1.0) - else: - face_normal = wp.mesh_eval_face_normal(mesh, query.face) - sign = query.sign + return query_result, sdf, sdf_grad, sdf_vel, closest_point, material_id - mesh_material_id = collider.face_material_index[global_face_id + query.face] - thickness = collider.material_thickness[mesh_material_id] - offset = x_local - cp - d = wp.length(offset) * sign - sdf = d - thickness +@wp.func +def collision_sdf( + x: wp.vec3, + environment_index: int, + collider: Collider, + body_q: wp.array[wp.transform], + body_qd: wp.array[wp.spatial_vector], + body_q_prev: wp.array[wp.transform], + dt: float, +): + min_sdf = float(_INFINITY) + sdf_grad = wp.vec3(0.0) + sdf_vel = wp.vec3(0.0) + closest_point = wp.vec3(0.0) + collider_id = int(_NULL_COLLIDER_ID) + material_id = int(0) # default material, always valid - if sdf < min_sdf: - min_sdf = sdf - if wp.abs(d) < _CLOSEST_POINT_NORMAL_EPSILON: - sdf_grad = face_normal - else: - sdf_grad = wp.normalize(offset) * sign + shared_worlds = environment_index == _ALL_COLLIDER_WORLDS + query_begin = int(0) + query_count = collider.collider_mesh.shape[0] - sdf_vel = wp.mesh_eval_velocity(mesh, query.face, query.u, query.v) - closest_point = cp - collider_id = m - material_id = mesh_material_id + if not shared_worlds: + query_begin = collider.world_collider_offsets[environment_index] + query_count = collider.world_collider_offsets[environment_index + 1] - query_begin - global_face_id += wp.mesh_get(mesh).indices.shape[0] // 3 + for query_offset in range(query_count): + stable_collider_id = query_offset + if not shared_worlds: + stable_collider_id = collider.world_collider_ids[query_begin + query_offset] + + query_result, sdf, query_grad, query_vel, query_closest_point, query_material_id = _query_collider_sdf( + x, collider, body_q, stable_collider_id + ) + if query_result and sdf < min_sdf: + min_sdf = sdf + sdf_grad = query_grad + sdf_vel = query_vel + closest_point = query_closest_point + collider_id = stable_collider_id + material_id = query_material_id # If closest collider has rigid motion, transform back to world frame # Do that as a second step to avoid requiring more registers inside bvh query loop @@ -265,6 +314,12 @@ def collision_sdf( return min_sdf, sdf_grad, sdf_vel, collider_id, material_id +@wp.func +def environment_from_offsets(index: int, offsets: wp.array[int]): + """Return the environment containing a packed-array index.""" + return wp.lower_bound(offsets, index + 1) - 1 + + @wp.kernel def collider_volumes_kernel( cell_volume: float, @@ -295,6 +350,7 @@ def project_outside_collider( velocity_gradients: wp.array[wp.mat33], particle_flags: wp.array[wp.int32], particle_mass: wp.array[float], + particle_environment: wp.array[int], collider: Collider, body_q: wp.array[wp.transform], body_qd: wp.array[wp.spatial_vector], @@ -318,6 +374,7 @@ def project_outside_collider( velocity_gradients: Current particle velocity gradients. particle_flags: Per-particle flags; particles without :attr:`ACTIVE` are skipped. particle_mass: Per-particle mass; zero-mass (kinematic) particles are skipped. + particle_environment: Per-particle world IDs, or null to query every collider. collider: Collider description and geometry. body_q: Rigid body transforms. body_qd: Rigid body velocities. @@ -339,9 +396,13 @@ def project_outside_collider( velocity_gradients_out[i] = vel_grad return + environment_index = int(_ALL_COLLIDER_WORLDS) + if particle_environment: + environment_index = particle_environment[i] + # project outside of collider sdf, sdf_gradient, sdf_vel, _collider_id, material_id = collision_sdf( - pos_adv, collider, body_q, body_qd, body_q_prev, dt + pos_adv, environment_index, collider, body_q, body_qd, body_q_prev, dt ) sdf_end = sdf - wp.dot(sdf_vel, sdf_gradient) * dt + collider.material_projection_threshold[material_id] @@ -374,6 +435,7 @@ def rasterize_collider_kernel( activation_distance: float, dt: float, node_positions: wp.array[wp.vec3], + node_environment_offsets: wp.array[int], node_volumes: wp.array[float], collider_sdf: wp.array[float], collider_velocity: wp.array[wp.vec3], @@ -399,6 +461,7 @@ def rasterize_collider_kernel( activation_distance: Distance (in voxels) below which to activate the collider. dt: Timestep length (used to scale adhesion and finite-difference velocity). node_positions: Grid node positions to sample at. + node_environment_offsets: Packed node offsets by world, or null to query every collider. node_volumes: Per-node integration volumes. collider_sdf: Output signed distance per node. collider_velocity: Output collider velocity per node. @@ -414,8 +477,11 @@ def rasterize_collider_kernel( bc_active = False sdf = _INFINITY else: + environment_index = int(_ALL_COLLIDER_WORLDS) + if node_environment_offsets: + environment_index = environment_from_offsets(i, node_environment_offsets) sdf, sdf_gradient, sdf_vel, collider_id, material_id = collision_sdf( - x, collider, body_q, body_qd, body_q_prev, dt + x, environment_index, collider, body_q, body_qd, body_q_prev, dt ) bc_active = sdf < activation_distance * voxel_size @@ -618,6 +684,7 @@ def rasterize_collider( collider_adhesion: wp.array[float], collider_ids: wp.array[int], temporary_store: fem.TemporaryStore, + node_environment_offsets: wp.array | None = None, ): """Rasterize collider signed-distance, normals, velocity, and material onto grid nodes. @@ -642,6 +709,8 @@ def rasterize_collider( collider_adhesion: Output adhesion per node [Pa]. collider_ids: Output collider index per node, or ``_NULL_COLLIDER_ID``. temporary_store: Temporary storage for intermediate buffers. + node_environment_offsets: Packed collision-node offsets by world. If ``None``, every node + queries every collider in stable order. """ collision_node_count = collider_position_field.dof_values.shape[0] @@ -670,6 +739,7 @@ def rasterize_collider( activation_distance, dt, collider_position_field.dof_values, + node_environment_offsets, collider_node_volume, collider_distance_field.dof_values, collider_velocity, diff --git a/newton/_src/solvers/implicit_mpm/render_grains.py b/newton/_src/solvers/implicit_mpm/render_grains.py index d1914a6099..1b0d86eaf5 100644 --- a/newton/_src/solvers/implicit_mpm/render_grains.py +++ b/newton/_src/solvers/implicit_mpm/render_grains.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + import warp as wp import warp.fem as fem @@ -46,6 +48,17 @@ def transform_grains( positions[pid, k] = p_pos_adv +@wp.kernel +def repeat_particle_environment( + particle_environment: wp.array[int], + grains_per_particle: int, + grain_environment: wp.array[int], +): + grain_index = wp.tid() + particle_index = grain_index // grains_per_particle + grain_environment[grain_index] = particle_environment[particle_index] + + @fem.integrand def advect_grains( s: fem.Sample, @@ -144,6 +157,9 @@ def update_render_grains( grains: wp.array, particle_radius: wp.array, dt: float, + *, + particle_environment: wp.array[int] | None = None, + temporary_store: fem.TemporaryStore | None = None, ): """Advect grain samples with the grid velocity and keep them inside the deformed particle. @@ -163,48 +179,76 @@ def update_render_grains( grains: 2D array of grain positions per particle to be updated in place. particle_radius: Per-particle radius used for projection. dt: Time step duration. + particle_environment: Per-particle world IDs for isolated multi-world advection. + temporary_store: Temporary storage used by grain binning and interpolation. """ - if state.velocity_field is None: + if getattr(state, "velocity_field", None) is None or grains.size == 0: return + grain_pos = grains.flatten() domain = fem.Cells(state.velocity_field.space.geometry) - grain_pic = fem.PicQuadrature(domain, positions=grain_pos) - - wp.launch( - advect_grains_from_particles, - dim=grains.shape, - inputs=[ - dt, - state_prev.particle_q, - state.particle_q, - state.mpm.particle_qd_grad, - grains, - ], - device=grains.device, - ) - - fem.interpolate( - advect_grains, - at=grain_pic, - values={ - "dt": dt, - "positions": grain_pos, - }, - fields={ - "grid_vel": state.velocity_field, - }, - device=grains.device, - ) - - wp.launch( - project_grains, - dim=grains.shape, - inputs=[ - particle_radius, - state.particle_q, - state.mpm.particle_transform, - grains, - ], - device=grains.device, - ) + grain_environment = None + try: + if particle_environment is not None: + grain_environment = fem.borrow_temporary( + temporary_store, + shape=grain_pos.shape, + dtype=int, + device=grains.device, + ) + wp.launch( + repeat_particle_environment, + dim=grain_pos.shape, + inputs=[particle_environment, grains.shape[1], grain_environment], + device=grains.device, + ) + + grain_pic = fem.PicQuadrature( + domain, + positions=grain_pos, + env_indices=grain_environment, + temporary_store=temporary_store, + ) + + wp.launch( + advect_grains_from_particles, + dim=grains.shape, + inputs=[ + dt, + state_prev.particle_q, + state.particle_q, + state.mpm.particle_qd_grad, + grains, + ], + device=grains.device, + ) + + fem.interpolate( + advect_grains, + at=grain_pic, + values={ + "dt": dt, + "positions": grain_pos, + }, + fields={ + "grid_vel": state.velocity_field, + }, + device=grains.device, + temporary_store=temporary_store, + ) + + wp.launch( + project_grains, + dim=grains.shape, + inputs=[ + particle_radius, + state.particle_q, + state.mpm.particle_transform, + grains, + ], + device=grains.device, + ) + finally: + if grain_environment is not None: + grain_environment.release() diff --git a/newton/_src/solvers/implicit_mpm/rheology_solver_kernels.py b/newton/_src/solvers/implicit_mpm/rheology_solver_kernels.py index 9d3f508720..9330d06470 100644 --- a/newton/_src/solvers/implicit_mpm/rheology_solver_kernels.py +++ b/newton/_src/solvers/implicit_mpm/rheology_solver_kernels.py @@ -308,6 +308,13 @@ def postprocess_stress_and_strain( """ tau_i = wp.tid() + strain_block_beg = strain_mat_offsets[tau_i] + strain_block_end = strain_mat_offsets[tau_i + 1] + if strain_block_beg == strain_block_end: + elastic_strain[tau_i] = vec6(0.0) + plastic_strain[tau_i] = vec6(0.0) + return + minus_elastic_strain = strain_rhs[tau_i] minus_elastic_strain -= unilateral_offset_to_strain_rhs(unilateral_strain_offset[tau_i]) comp_block_beg = compliance_mat_offsets[tau_i] @@ -317,9 +324,7 @@ def postprocess_stress_and_strain( minus_elastic_strain += compliance_mat_values[b] * stress[sig_i] world_plastic_strain = minus_elastic_strain - block_beg = strain_mat_offsets[tau_i] - block_end = strain_mat_offsets[tau_i + 1] - for b in range(block_beg, block_end): + for b in range(strain_block_beg, strain_block_end): u_i = strain_mat_columns[b] world_plastic_strain += _symmetric_part_op(strain_mat_values[b], velocity[u_i]) @@ -705,7 +710,6 @@ def apply_stress_delta_impl( """Updates particle velocities from a local stress delta.""" block_beg = strain_mat_offsets[tau_i] - if wp.static(strain_velocity_node_count > 0): for bk in range(strain_velocity_node_count): b = block_beg + bk @@ -941,6 +945,10 @@ def jacobi_solve_kernel_impl( ): tau_i = wp.tid() + if strain_mat_offsets[tau_i] == strain_mat_offsets[tau_i + 1]: + delta_correction[tau_i] = vec6(0.0) + return + local_strain = wp.static(make_compute_local_strain(has_compliance_mat, strain_velocity_node_count))( tau_i, compliance_mat_offsets, @@ -1012,6 +1020,10 @@ def gs_solve_kernel_impl( for color_offset in range(color_beg, color_end, launch_dim): beg, end = color_blocks[0, color_offset], color_blocks[1, color_offset] for tau_i in range(beg, end): + if strain_mat_offsets[tau_i] == strain_mat_offsets[tau_i + 1]: + delta_correction[tau_i] = vec6(0.0) + continue + local_strain = wp.static(make_compute_local_strain(has_compliance_mat, strain_velocity_node_count))( tau_i, compliance_mat_offsets, diff --git a/newton/_src/solvers/implicit_mpm/solve_rheology.py b/newton/_src/solvers/implicit_mpm/solve_rheology.py index 11cfbb07ae..4d539fa44e 100644 --- a/newton/_src/solvers/implicit_mpm/solve_rheology.py +++ b/newton/_src/solvers/implicit_mpm/solve_rheology.py @@ -1,12 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + import gc import math from collections.abc import Sequence from dataclasses import dataclass from typing import Any +import numpy as np import warp as wp import warp.fem as fem import warp.sparse as sp @@ -53,6 +56,7 @@ ) _TILED_SUM_BLOCK_DIM = 512 +_STRESS_DOF_COUNT = 6 @wp.kernel @@ -68,12 +72,98 @@ def _tiled_sum_kernel( wp.tile_store(partial_sums[1], wp.tile_max(tile), offset=block_id) +@wp.kernel +def _batched_sum_max_kernel( + data: wp.array2d[float], + batch_offsets: wp.array[int], + result: wp.array2d[float], +): + batch, lane = wp.tid() + + batch_start = batch_offsets[batch] + lane + batch_end = batch_offsets[batch + 1] + + sum_value = float(0.0) + max_value = float(0.0) + for i in range(batch_start, batch_end, wp.block_dim()): + sum_value += data[0, i] + max_value = wp.max(max_value, data[1, i]) + + wp.tile_store(result[0], wp.tile_sum(wp.tile(sum_value)), offset=batch) + wp.tile_store(result[1], wp.tile_max(wp.tile(max_value)), offset=batch) + + +@wp.kernel +def _scale_offsets_kernel( + offsets: wp.array[int], + scale: int, + scaled_offsets: wp.array[int], +): + i = wp.tid() + scaled_offsets[i] = scale * offsets[i] + + +@wp.kernel +def _compute_environment_l2_tolerance_scales( + environment_offsets: wp.array[int], + tolerance_scales: wp.array[float], +): + environment = wp.tid() + environment_size = environment_offsets[environment + 1] - environment_offsets[environment] + tolerance_scales[environment] = wp.sqrt(float(1 + environment_size)) + + +@wp.kernel +def _scale_linear_system_by_environment( + environment_offsets: wp.array[int], + tolerance_scales: wp.array[float], + inverse: bool, + rhs: wp.array[vec6], + solution: wp.array[vec6], +): + environment, lane = wp.tid() + scale = tolerance_scales[environment] + factor = wp.where(inverse, 1.0 / scale, scale) + begin = environment_offsets[environment] + lane + end = environment_offsets[environment + 1] + for i in range(begin, end, wp.block_dim()): + rhs[i] = factor * rhs[i] + solution[i] = factor * solution[i] + + class ArraySquaredNorm: """Utility to compute squared L2 norm of a large array via tiled reductions.""" - def __init__(self, max_length: int, device=None, temporary_store=None): + def __init__(self, max_length: int, batch_offsets: wp.array[int] | None = None, device=None, temporary_store=None): self.tile_size = _TILED_SUM_BLOCK_DIM self.device = device + self.batch_offsets = batch_offsets + + self.partial_sums_a = None + self.partial_sums_b = None + self.batch_result = None + self.sum_launch = None + self.batch_sum_launch = None + + if batch_offsets is not None: + if not self.device.is_cuda: + self.tile_size = 1 + + batch_count = batch_offsets.shape[0] - 1 + self.batch_result = fem.borrow_temporary( + temporary_store, shape=(2, batch_count), dtype=float, device=self.device + ) + self.batch_result.zero_() + self.batch_sum_launch = wp.launch( + _batched_sum_max_kernel, + dim=(batch_count, self.tile_size), + inputs=(self.batch_result, batch_offsets), + outputs=(self.batch_result,), + block_dim=self.tile_size, + device=self.device, + record_cmd=True, + ) + return num_blocks = (max_length + self.tile_size - 1) // self.tile_size self.partial_sums_a = fem.borrow_temporary( @@ -91,6 +181,7 @@ def __init__(self, max_length: int, device=None, temporary_store=None): inputs=(self.partial_sums_a,), outputs=(self.partial_sums_b,), block_dim=self.tile_size, + device=self.device, record_cmd=True, ) @@ -106,6 +197,11 @@ def compute_squared_norm(self, data: wp.array[Any]): device=data.device, ) + if self.batch_sum_launch is not None: + self.batch_sum_launch.set_param_at_index(0, data) + self.batch_sum_launch.launch() + return self.batch_result + array_length = data.shape[1] flip_flop = False @@ -130,7 +226,7 @@ def compute_squared_norm(self, data: wp.array[Any]): def release(self): """Return borrowed temporaries to their pool.""" - for attr in ("partial_sums_a", "partial_sums_b"): + for attr in ("partial_sums_a", "partial_sums_b", "batch_result"): temporary = getattr(self, attr, None) if temporary is not None: temporary.release() @@ -151,9 +247,35 @@ def update_condition( condition: wp.array[int], ): cur_it = iteration[0] + solve_granularity - stop = ( - residual[0, 0] < residual_threshold * l2_scale and residual[1, 0] < residual_threshold - ) or cur_it > max_iterations + converged = bool(True) + for batch in range(residual.shape[1]): + converged = converged and residual[0, batch] < residual_threshold * l2_scale + converged = converged and residual[1, batch] < residual_threshold + + stop = converged or cur_it > max_iterations + + iteration[0] = cur_it + condition[0] = wp.where(stop, 0, 1) + + +@wp.kernel +def update_batched_condition( + residual_threshold: float, + l2_tolerance_scales: wp.array[float], + solve_granularity: int, + max_iterations: int, + residual: wp.array2d[float], + iteration: wp.array[int], + condition: wp.array[int], +): + cur_it = iteration[0] + solve_granularity + converged = bool(True) + for batch in range(residual.shape[1]): + scale = l2_tolerance_scales[batch] + converged = converged and residual[0, batch] < residual_threshold * scale * scale + converged = converged and residual[1, batch] < residual_threshold + + stop = converged or cur_it > max_iterations iteration[0] = cur_it condition[0] = wp.where(stop, 0, 1) @@ -239,6 +361,9 @@ class RheologyData: node, shape ``[strain_count, 6]``. stress: In/out stress per strain node (rotated internally), shape ``[strain_count, 6]``. + strain_environment_offsets: Strain-node offsets delimiting independent + environments, shape ``[environment_count + 1]``. ``None`` for a + single shared solve. """ strain_mat: sp.BsrMatrix @@ -254,6 +379,7 @@ class RheologyData: elastic_strain_delta: wp.array[vec6] plastic_strain_delta: wp.array[vec6] stress: wp.array[vec6] + strain_environment_offsets: wp.array[int] | None = None has_viscosity: bool = False has_dilatancy: bool = False @@ -533,6 +659,7 @@ def __init__( # Utility to compute the squared norm of the residual self._residual_squared_norm_computer = ArraySquaredNorm( max_length=self.size, + batch_offsets=self.rheology.strain_environment_offsets, device=self.device, temporary_store=temporary_store, ) @@ -1312,9 +1439,36 @@ def __init__( dtype = self.rheology.compliance_mat.dtype device = self.rheology.compliance_mat.device - self.linear_operator = LinearOperator(shape=shape, dtype=dtype, device=device, matvec=self._delassus_matvec) + self._batch_offsets = None + if self.rheology.strain_environment_offsets is not None: + strain_environment_offsets = self.rheology.strain_environment_offsets + self._batch_offsets = fem.borrow_temporary( + temporary_store, + shape=strain_environment_offsets.shape, + dtype=int, + device=device, + ) + wp.launch( + _scale_offsets_kernel, + dim=strain_environment_offsets.shape[0], + inputs=[strain_environment_offsets, _STRESS_DOF_COUNT], + outputs=[self._batch_offsets], + device=device, + ) + + self.linear_operator = LinearOperator( + shape=shape, + dtype=dtype, + device=device, + matvec=self._delassus_matvec, + batch_offsets=self._batch_offsets, + ) self.preconditioner = LinearOperator( - shape=shape, dtype=dtype, device=device, matvec=self._preconditioner_matvec + shape=shape, + dtype=dtype, + device=device, + matvec=self._preconditioner_matvec, + batch_offsets=self._batch_offsets, ) def _delassus_matvec(self, x: wp.array[vec6], y: wp.array[vec6], z: wp.array[vec6], alpha: float, beta: float): @@ -1342,7 +1496,31 @@ def _preconditioner_matvec(self, x, y, z, alpha, beta): ], ) - def solve(self, tol: float, tolerance_scale: float, max_iterations: int, use_graph: bool, verbose: bool): + def _scale_batched_system(self, tolerance_scales: wp.array[float], inverse: bool): + environment_offsets = self.rheology.strain_environment_offsets + block_dim = 256 if self.momentum.velocity.device.is_cuda else 1 + wp.launch( + _scale_linear_system_by_environment, + dim=(tolerance_scales.shape[0], block_dim), + inputs=[ + environment_offsets, + tolerance_scales, + inverse, + self.rheology.plastic_strain_delta, + self.rheology.stress, + ], + block_dim=block_dim, + device=self.momentum.velocity.device, + ) + + def solve( + self, + tol: float, + tolerance_scale: float | wp.array, + max_iterations: int, + use_graph: bool, + verbose: bool, + ): self.delassus_operator.apply_velocity_delta( self.momentum.velocity, self.rheology.elastic_strain_delta, @@ -1351,29 +1529,37 @@ def solve(self, tol: float, tolerance_scale: float, max_iterations: int, use_gra beta=-1.0, ) - with _ScopedDisableGC(): - end_iter, residual, _ = self._method_fn( - A=self.linear_operator, - M=self.preconditioner, - b=self.rheology.plastic_strain_delta, - x=self.rheology.stress, - atol=tol * tolerance_scale, - tol=tol, - maxiter=max_iterations, - check_every=0 if use_graph else 10, - use_cuda_graph=use_graph, - ) + is_batched = self._batch_offsets is not None + if is_batched: + self._scale_batched_system(tolerance_scale, inverse=True) + + try: + with _ScopedDisableGC(): + end_iter, residual, atol = self._method_fn( + A=self.linear_operator, + M=self.preconditioner, + b=self.rheology.plastic_strain_delta, + x=self.rheology.stress, + atol=tol if is_batched else tol * tolerance_scale, + tol=tol, + maxiter=max_iterations, + check_every=0 if use_graph else 10, + use_cuda_graph=use_graph, + ) + finally: + if is_batched: + self._scale_batched_system(tolerance_scale, inverse=False) # With use_cuda_graph=True the solver returns end_iter and residual as - # length-1 device arrays so the caller need not synchronize. Read them - # back only for the verbose report, and never while an outer capture is - # recording: a device-to-host copy there serializes the capturing stream - # (CUDA error 906). + # device arrays so the caller need not synchronize. Read them back only + # for the verbose report, and never while an outer capture is recording: + # a device-to-host copy there serializes the capturing stream (CUDA + # error 906). Batched solves report the largest residual. if verbose and not (use_graph and self.momentum.velocity.device.is_capturing): if use_graph: end_iter = end_iter.numpy()[0] - residual = residual.numpy()[0] - res = math.sqrt(residual) / tolerance_scale + residual, _ = _linear_solver_result_norms(residual, atol, use_graph) + res = residual if is_batched else residual / tolerance_scale print(f"{self.name} terminated after {end_iter} iterations with residual {res}") @property @@ -1382,6 +1568,16 @@ def name(self): def release(self): self.delta_velocity.release() + if self._batch_offsets is not None: + self._batch_offsets.release() + self._batch_offsets = None + + +def _linear_solver_result_norms(residual, atol, use_graph: bool) -> tuple[float, float]: + if use_graph: + residual = math.sqrt(float(residual.numpy().max())) + atol = math.sqrt(float(atol.numpy().max())) + return residual, atol class _ContactSolver: @@ -1574,12 +1770,21 @@ def release(self): super().release() +def _nonlinear_solver_result_norms(residual, l2_tolerance_scale: float | np.ndarray) -> tuple[float, float]: + """Return the largest independently scaled residual across solver batches.""" + + return ( + float(np.max(np.sqrt(residual[0]) / l2_tolerance_scale)), + math.sqrt(float(residual[1].max())), + ) + + def _run_solver_loop( rheology_solver: _RheologySolver, contact_solver: _ContactSolver, max_iterations: int, tolerance: float, - l2_tolerance_scale: float, + l2_tolerance_scale: float | wp.array, use_graph: bool, verbose: bool, temporary_store: fem.TemporaryStore, @@ -1599,19 +1804,34 @@ def do_iteration_with_condition(): contact_solver.solve() rheology_solver.solve() residual = rheology_solver.eval_residual() - wp.launch( - update_condition, - dim=1, - inputs=[ - tolerance * tolerance, - l2_tolerance_scale * l2_tolerance_scale, - solve_granularity, - max_iterations, - residual, - iteration, - condition, - ], - ) + if rheology_solver.rheology.strain_environment_offsets is None: + wp.launch( + update_condition, + dim=1, + inputs=[ + tolerance * tolerance, + l2_tolerance_scale * l2_tolerance_scale, + solve_granularity, + max_iterations, + residual, + iteration, + condition, + ], + ) + else: + wp.launch( + update_batched_condition, + dim=1, + inputs=[ + tolerance * tolerance, + l2_tolerance_scale, + solve_granularity, + max_iterations, + residual, + iteration, + condition, + ], + ) device = rheology_solver.device if device.is_capturing: @@ -1626,7 +1846,12 @@ def do_iteration_with_condition(): if verbose: residual = rheology_solver.eval_residual().numpy() - res_l2, res_linf = math.sqrt(residual[0, 0]) / l2_tolerance_scale, math.sqrt(residual[1, 0]) + host_tolerance_scale = ( + l2_tolerance_scale + if rheology_solver.rheology.strain_environment_offsets is None + else l2_tolerance_scale.numpy() + ) + res_l2, res_linf = _nonlinear_solver_result_norms(residual, host_tolerance_scale) print( f"{rheology_solver.name} terminated after {iteration_and_condition.numpy()[0]} iterations with residuals {res_l2}, {res_linf}" ) @@ -1634,6 +1859,11 @@ def do_iteration_with_condition(): iteration_and_condition.release() else: solve_granularity = rheology_solver.solve_granularity + host_tolerance_scale = ( + l2_tolerance_scale + if rheology_solver.rheology.strain_environment_offsets is None + else l2_tolerance_scale.numpy() + ) for batch in range(max_iterations // solve_granularity): for _k in range(solve_granularity): @@ -1641,7 +1871,7 @@ def do_iteration_with_condition(): rheology_solver.solve() residual = rheology_solver.eval_residual().numpy() - res_l2, res_linf = math.sqrt(residual[0, 0]) / l2_tolerance_scale, math.sqrt(residual[1, 0]) + res_l2, res_linf = _nonlinear_solver_result_norms(residual, host_tolerance_scale) if verbose: print( @@ -1722,7 +1952,6 @@ def solve_rheology( """ verbose = verbose if verbose is not None else wp.config.log_level <= wp.LOG_DEBUG - subgrid_collisions = collision.collider_mat.nnz > 0 if subgrid_collisions: contact_solver = _SubgridContactSolver(momentum, collision, temporary_store) @@ -1732,7 +1961,25 @@ def solve_rheology( contact_solver.apply_initial_guess() delassus_operator = _DelassusOperator(rheology, momentum, temporary_store) - tolerance_scale = math.sqrt(1 + delassus_operator.size) + batch_tolerance_scales = None + if rheology.strain_environment_offsets is None: + tolerance_scale = math.sqrt(1 + delassus_operator.size) + else: + environment_count = rheology.strain_environment_offsets.shape[0] - 1 + batch_tolerance_scales = fem.borrow_temporary( + temporary_store, + shape=(environment_count,), + dtype=float, + device=momentum.velocity.device, + ) + wp.launch( + _compute_environment_l2_tolerance_scales, + dim=environment_count, + inputs=[rheology.strain_environment_offsets], + outputs=[batch_tolerance_scales], + device=momentum.velocity.device, + ) + tolerance_scale = batch_tolerance_scales solvers = (solver,) if isinstance(solver, str) else tuple(solver) if len(solvers) == 0: @@ -1755,6 +2002,8 @@ def solve_rheology( delassus_operator.postprocess_stress_and_strain() delassus_operator.release() contact_solver.release() + if batch_tolerance_scales is not None: + batch_tolerance_scales.release() return None # linear solver as warmstart @@ -1787,12 +2036,21 @@ def solve_rheology( rheology_solver.apply_initial_guess() solve_graph = _run_solver_loop( - rheology_solver, contact_solver, max_iterations, tolerance, tolerance_scale, use_graph, verbose, temporary_store + rheology_solver, + contact_solver, + max_iterations, + tolerance, + tolerance_scale, + use_graph, + verbose, + temporary_store, ) # release temporary storage rheology_solver.release() contact_solver.release() + if batch_tolerance_scales is not None: + batch_tolerance_scales.release() delassus_operator.postprocess_stress_and_strain() delassus_operator.release() diff --git a/newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py b/newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py index efb8ceca60..ae489d98da 100644 --- a/newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py +++ b/newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py @@ -5,8 +5,10 @@ from __future__ import annotations +import operator from collections.abc import Sequence from dataclasses import dataclass +from itertools import pairwise from typing import Literal import numpy as np @@ -38,9 +40,11 @@ EPSILON, INFINITY, YIELD_PARAM_LENGTH, + _rebuild_capacity, advect_particles, allocate_by_voxels, average_elastic_parameters, + build_active_particle_mask, collision_weight_field, compliance_form, compute_bounds, @@ -66,12 +70,18 @@ make_inverse_rotate_vectors, make_rotate_vectors, mark_active_cells, + mark_active_cells_by_environment, mass_form, mat11, mat13, mat31, mat66, node_color, + record_volume_rebuild_status, + reset_mpm_collider_history, + reset_mpm_grid_warmstart, + reset_mpm_particle_history, + reset_mpm_point_warmstart, rotate_matrix_columns, rotate_matrix_rows, scatter_field_dof_values, @@ -79,6 +89,7 @@ strain_rhs, update_particle_frames, update_particle_strains, + voxel_coordinates, ) @@ -94,6 +105,54 @@ def _as_2d_array(array, shape, dtype): ) +def _sparse_grid_rebuild_error(status: int) -> RuntimeError: + capacity_flags = ( + (wp.Volume.REBUILD_VOXEL_CAPACITY_EXCEEDED, "active voxels"), + (wp.Volume.REBUILD_LEAF_CAPACITY_EXCEEDED, "leaf nodes"), + (wp.Volume.REBUILD_LOWER_CAPACITY_EXCEEDED, "lower internal nodes"), + (wp.Volume.REBUILD_UPPER_CAPACITY_EXCEEDED, "upper internal nodes"), + ) + exceeded = [name for flag, name in capacity_flags if status & flag] + if exceeded: + details = ", ".join(exceeded) + suggestions = [] + if status & wp.Volume.REBUILD_VOXEL_CAPACITY_EXCEEDED: + suggestions.append("increase Config.max_active_cell_count") + if status & wp.Volume.REBUILD_LEAF_CAPACITY_EXCEEDED: + suggestions.append("set or increase Config.max_leaf_node_count") + if status & wp.Volume.REBUILD_LOWER_CAPACITY_EXCEEDED: + suggestions.append("set or increase Config.max_lower_node_count") + if status & wp.Volume.REBUILD_UPPER_CAPACITY_EXCEEDED: + suggestions.append("set or increase Config.max_upper_node_count") + if status & ( + wp.Volume.REBUILD_LEAF_CAPACITY_EXCEEDED + | wp.Volume.REBUILD_LOWER_CAPACITY_EXCEEDED + | wp.Volume.REBUILD_UPPER_CAPACITY_EXCEEDED + ): + suggestions.append("reduce the active grid's spatial spread") + suggestion = " or ".join(suggestions) + return RuntimeError( + f"Implicit MPM sparse grid rebuild capacity was exceeded for {details} (status {status}). " + f"To avoid overflow, {suggestion}." + ) + return RuntimeError(f"Implicit MPM sparse grid rebuild failed with status {status}.") + + +def _validate_sparse_grid_node_capacity(name: str, value: int) -> int: + """Validate one optional NanoVDB hierarchy capacity.""" + if isinstance(value, bool): + raise ValueError(f"Config.{name} must be -1 or a positive integer, got {value!r}.") + try: + capacity = operator.index(value) + except TypeError as error: + raise ValueError(f"Config.{name} must be -1 or a positive integer, got {value!r}.") from error + if capacity != -1 and not 0 < capacity <= np.iinfo(np.uint32).max: + raise ValueError( + f"Config.{name} must be -1 or a positive integer no greater than {np.iinfo(np.uint32).max}, got {capacity}." + ) + return capacity + + def _make_grid_basis_space(grid: fem.Geometry, basis_str: str, family: fem.Polynomial | None = None): assert len(basis_str) >= 2 @@ -212,6 +271,10 @@ def __init__(self): self.collider_total_volumes = None self.collider_node_volume = None + self.velocity_environment_offsets = None + self.collider_environment_offsets = None + self.strain_environment_offsets = None + def rebuild_function_spaces( self, pic: fem.PicQuadrature, @@ -219,6 +282,7 @@ def rebuild_function_spaces( strain_basis_str: str, collider_basis_str: str, max_cell_count: int, + environment_first: bool, temporary_store: fem.TemporaryStore, ): """Define velocity and strain function spaces over the given geometry.""" @@ -228,6 +292,8 @@ def rebuild_function_spaces( use_pic_collider_basis = collider_basis_str[:3] == "pic" use_pic_strain_basis = strain_basis_str[:3] == "pic" + # The rebuildable sparse grid (like the fixed grid) reuses the same geometry object + # across steps, so refresh retained topologies before rebuilding their partitions. if self.domain.geometry is not self.grid: self.grid = self.domain.geometry @@ -242,6 +308,15 @@ def rebuild_function_spaces( self._collision_basis = _make_grid_basis_space( self.grid, collider_basis_str, family=fem.Polynomial.EQUISPACED_CLOSED ) + else: + topologies = {id(self._velocity_basis.topology): self._velocity_basis.topology} + if not use_pic_strain_basis: + topologies[id(self._strain_basis.topology)] = self._strain_basis.topology + if not use_pic_collider_basis: + topologies[id(self._collision_basis.topology)] = self._collision_basis.topology + for topology in topologies.values(): + if hasattr(topology, "rebuild"): + topology.rebuild() # Point-based basis space needs to be rebuilt even when the geo does not change if use_pic_strain_basis: @@ -249,11 +324,13 @@ def rebuild_function_spaces( if use_pic_collider_basis: self._collision_basis = _make_pic_basis_space(pic, collider_basis_str) - self._create_velocity_function_space(temporary_store, max_cell_count) - self._create_collider_function_space(temporary_store, max_cell_count) - self._create_strain_function_space(temporary_store, max_cell_count) + self._create_velocity_function_space(temporary_store, max_cell_count, environment_first) + self._create_collider_function_space(temporary_store, max_cell_count, environment_first) + self._create_strain_function_space(temporary_store, max_cell_count, environment_first) - def _create_velocity_function_space(self, temporary_store: fem.TemporaryStore, max_cell_count: int = -1): + def _create_velocity_function_space( + self, temporary_store: fem.TemporaryStore, max_cell_count: int, environment_first: bool + ): """Create velocity and fraction spaces and their partition/restriction.""" domain = self.domain @@ -269,6 +346,7 @@ def _create_velocity_function_space(self, temporary_store: fem.TemporaryStore, m geometry_partition=domain.geometry_partition, with_halo=False, max_node_count=max_vel_node_count, + environment_first=environment_first, temporary_store=temporary_store, ) vel_space_restriction = fem.make_space_restriction( @@ -277,13 +355,17 @@ def _create_velocity_function_space(self, temporary_store: fem.TemporaryStore, m self._velocity_space = velocity_space self._vel_space_restriction = vel_space_restriction + self.velocity_environment_offsets = vel_space_partition.env_offsets if environment_first else None - def _create_collider_function_space(self, temporary_store: fem.TemporaryStore, max_cell_count: int = -1): + def _create_collider_function_space( + self, temporary_store: fem.TemporaryStore, max_cell_count: int, environment_first: bool + ): """Create collider function space and its partition/restriction.""" if self._velocity_basis == self._collision_basis: self._collision_space = self._velocity_space self._collision_space_restriction = self._vel_space_restriction + self.collider_environment_offsets = self.velocity_environment_offsets return domain = self.domain @@ -303,6 +385,7 @@ def _create_collider_function_space(self, temporary_store: fem.TemporaryStore, m geometry_partition=domain.geometry_partition, with_halo=False, max_node_count=max_collision_node_count, + environment_first=environment_first, temporary_store=temporary_store, ) collision_space_restriction = fem.make_space_restriction( @@ -311,8 +394,11 @@ def _create_collider_function_space(self, temporary_store: fem.TemporaryStore, m self._collision_space = collision_space self._collision_space_restriction = collision_space_restriction + self.collider_environment_offsets = collision_space_partition.env_offsets if environment_first else None - def _create_strain_function_space(self, temporary_store: fem.TemporaryStore, max_cell_count: int = -1): + def _create_strain_function_space( + self, temporary_store: fem.TemporaryStore, max_cell_count: int, environment_first: bool + ): """Create symmetric strain space (P0 or Q1) and its partition/restriction.""" domain = self.domain @@ -330,6 +416,7 @@ def _create_strain_function_space(self, temporary_store: fem.TemporaryStore, max geometry_partition=domain.geometry_partition, with_halo=False, max_node_count=max_strain_node_count, + environment_first=environment_first, temporary_store=temporary_store, ) @@ -339,6 +426,7 @@ def _create_strain_function_space(self, temporary_store: fem.TemporaryStore, max self._sym_strain_space = sym_strain_space self._strain_space_restriction = strain_space_restriction + self.strain_environment_offsets = strain_space_partition.env_offsets if environment_first else None def require_velocity_space_fields(self, has_compliant_particles: bool): velocity_basis = self._velocity_basis @@ -610,10 +698,26 @@ def require_collider_previous_position(self, collider_body_q: wp.array | None): elif self.body_q_prev is None or self.body_q_prev.shape != collider_body_q.shape: self.body_q_prev = wp.clone(collider_body_q) - def save_collider_current_position(self, collider_body_q: wp.array | None): + def save_collider_current_position( + self, + collider_body_q: wp.array[wp.transform] | None, + body_world: wp.array[wp.int32] | None = None, + world_mask: wp.array[wp.bool] | None = None, + ): + had_previous = self.body_q_prev is not None and ( + collider_body_q is None or self.body_q_prev.shape == collider_body_q.shape + ) self.require_collider_previous_position(collider_body_q) if collider_body_q is not None: - self.body_q_prev.assign(collider_body_q) + if world_mask is None or body_world is None or not had_previous: + self.body_q_prev.assign(collider_body_q) + else: + wp.launch( + reset_mpm_collider_history, + dim=collider_body_q.shape[0], + inputs=[body_world, world_mask, collider_body_q, self.body_q_prev], + device=collider_body_q.device, + ) class SolverImplicitMPM(SolverBase, CouplingInterface): @@ -633,6 +737,29 @@ class SolverImplicitMPM(SolverBase, CouplingInterface): parameters and state variables (e.g. ``mpm:young_modulus``, ``mpm:friction``, ``mpm:particle_elastic_strain``). + Multi-world models retain the shared FEM topology by default. Set + :attr:`Config.separate_worlds` to use an independent FEM environment for + each world. In isolated mode, every MPM particle must belong to a local + world and particles must be stored contiguously by world. World-local + colliders affect only that world; global static colliders and global + colliders backed by kinematic bodies affect every world, while global + colliders backed by dynamic bodies are rejected. + + A sparse grid is rebuildable when :attr:`Config.max_active_cell_count` is + positive, :attr:`Config.grid_padding` is zero, the velocity basis is + ``"Q1"``, and the strain and collider bases support rebuilding. Cell and + node capacities are totals across all FEM environments, and resolved + capacities must satisfy ``upper <= lower <= leaf <= active``. + + Outer graph capture requires CUDA, an enabled memory pool, conditional + graph support, ``enable_timers=False``, positive active-cell capacity, and + either a fixed or rebuildable sparse grid. Construction creates persistent + capture resources internally. Run an uncaptured warm-up step for rheology + modes that build an inner graph lazily, and keep model topology, solver + configuration, captured buffers, and step arguments fixed across replays. + Keep :meth:`reset` outside capture and call :meth:`check_status` after + replay to report sparse-grid capacity failures. + [1] https://doi.org/10.1145/2897824.2925877 Args: @@ -642,7 +769,14 @@ class SolverImplicitMPM(SolverBase, CouplingInterface): allocations across steps. verbose: If True, enable verbose solver output. If False, suppress details. If None, enable verbose output when ``wp.config.log_level`` is configured for debug logging. - enable_timers: Enable per-section wall-clock timings. + enable_timers: Enable per-section wall-clock timings. This must be + ``False`` when :meth:`step` is recorded in an outer graph + because section timers synchronize the device. + + Raises: + ValueError: If an isolated multi-world model contains global MPM + particles, particle world IDs outside ``[-1, model.world_count)``, + or particles that are not stored contiguously by world. """ @dataclass @@ -668,7 +802,13 @@ class Config: warmstart solvers left-to-right, e.g. ``("cr", "gs")`` or ``("cg", "jacobi", "gs")``.""" warmstart_mode: Literal["none", "auto", "particles", "grid", "smoothed"] = "auto" - """Warmstart mode to use for the rheology solver.""" + """Warmstart mode to use for the rheology solver. + + ``"auto"`` uses particle-backed stress for rebuildable sparse grids + and ``P1d``/``Q1d`` strain bases, and grid-backed stress otherwise. + Grid-backed ``"grid"`` and ``"smoothed"`` modes are not supported for + rebuildable sparse grids because their topology changes in place. + """ collider_velocity_mode: Literal["forward", "backward"] = "forward" """Collider velocity computation mode. ``'forward'`` uses the current velocity, ``'backward'`` uses the previous timestep position.""" @@ -677,11 +817,50 @@ class Config: voxel_size: float = 0.1 """Size of the grid voxels.""" grid_type: Literal["sparse", "dense", "fixed"] = "sparse" - """Type of grid to use.""" + """Type of grid to use. + + A capacity-bounded ``"sparse"`` grid is rebuilt in place. Dense grids + read dynamic bounds on the host. + """ grid_padding: int = 0 """Number of empty cells to add around particles when allocating the grid.""" max_active_cell_count: int = -1 - """Maximum number of active cells to use for active subsets of dense grids. -1 means unlimited.""" + """Maximum number of active grid cells across all worlds. + + A positive value reserves persistent sparse-grid capacity and bounds + active subsets of dense and fixed grids. + ``-1`` retains per-step sparse-grid allocation and uses exact active + counts elsewhere. Call :meth:`check_status` after graph replay to + detect sparse-grid overflow. + """ + max_leaf_node_count: int = -1 + """Maximum NanoVDB leaf-node count across all worlds. + + This independently bounds leaf topology for a rebuildable sparse grid. + ``-1`` reserves one leaf per :attr:`max_active_cell_count`, which is the + worst case for arbitrarily scattered active cells. Set an explicit + value only when application-level spatial bounds guarantee a tighter + limit. All node-capacity settings are validated at construction but + used only for rebuildable sparse grids. + """ + max_lower_node_count: int = -1 + """Maximum NanoVDB lower internal-node count across all worlds. + + ``-1`` estimates the initial packed topology and reserves 16 times its + lower-node count, capped by the resolved leaf-node capacity. An + explicit value budgets spatial-spread headroom independently from + active cells. Only used for rebuildable sparse grids. + """ + max_upper_node_count: int = -1 + """Maximum NanoVDB upper internal-node count across all worlds. + + ``-1`` estimates the initial packed topology and reserves 16 times its + upper-node count, capped by the resolved lower-node capacity. Upper + nodes cover regions 4096 voxels wide and are substantially larger than + lower or leaf nodes, so applications with known spatial bounds can use + this field to budget them explicitly. Only used for rebuildable sparse + grids. + """ transfer_scheme: Literal["apic", "pic"] = "apic" """Transfer scheme to use for particle-grid transfers.""" integration_scheme: Literal["pic", "gimp"] = "pic" @@ -711,6 +890,22 @@ class Config: """Velocity basis function. Common values are ``"Q1"``, ``"B2"``, or ``"B3"``.""" + separate_worlds: bool = False + """Use independent FEM environments for each world in a multi-world model. + + The default ``False`` retains the legacy shared-grid behavior. Set to + ``True`` to isolate grid mass, momentum, stress, and collider response + by world. Isolated multi-world models require every MPM particle to + belong to a local world and particles to be stored contiguously by + world. See :ref:`implicit-mpm-worlds` for how this mode interprets + world assignment and collider ownership. + + .. experimental:: + + Isolated multi-world MPM configuration and behavior may change + without prior notice. + """ + @classmethod def register_custom_attributes(cls, builder: newton.ModelBuilder) -> None: """Register MPM-specific custom attributes in the 'mpm' namespace. @@ -926,6 +1121,44 @@ def __init__( ): super().__init__(model) + self._initial_particle_count = int(model.particle_count) + self._initial_world_count = int(model.world_count) + + if self._initial_particle_count == 0: + raise ValueError("SolverImplicitMPM requires at least one particle.") + + self._separate_worlds = bool(config.separate_worlds and self._initial_world_count > 1) + self._environment_count = self._initial_world_count if self._separate_worlds else 1 + self._particle_environment = model.particle_world if self._separate_worlds else None + self._particle_world_ranges = None + + if self._separate_worlds: + particle_world = model.particle_world.numpy() + if np.any(particle_world < -1) or np.any(particle_world >= model.world_count): + raise ValueError( + "SolverImplicitMPM found invalid MPM particle world IDs; expected values in " + "[-1, model.world_count)." + ) + if np.any(particle_world == -1): + raise ValueError( + "SolverImplicitMPM cannot isolate a multi-world model containing global MPM particles; " + "replicate the particles into each world or set Config.separate_worlds=False for legacy " + "coupled behavior." + ) + + world_boundaries = np.searchsorted(particle_world, np.arange(model.world_count + 1)) + particle_world_ranges = tuple((int(begin), int(end)) for begin, end in pairwise(world_boundaries)) + ranges_cover_particles = world_boundaries[0] == 0 and world_boundaries[-1] == particle_world.shape[0] + ranges_match_worlds = all( + np.all(particle_world[begin:end] == world) for world, (begin, end) in enumerate(particle_world_ranges) + ) + if not ranges_cover_particles or not ranges_match_worlds: + raise ValueError( + "SolverImplicitMPM requires MPM particles to be stored contiguously by world when " + "Config.separate_worlds=True." + ) + self._particle_world_ranges = particle_world_ranges + self._mpm_model = ImplicitMPMModel(model, config) self.max_iterations = config.max_iterations @@ -941,12 +1174,35 @@ def __init__( self.grid_padding = config.grid_padding self.grid_type = config.grid_type + self.max_active_cell_count = config.max_active_cell_count + self.max_leaf_node_count = _validate_sparse_grid_node_capacity( + "max_leaf_node_count", config.max_leaf_node_count + ) + self.max_lower_node_count = _validate_sparse_grid_node_capacity( + "max_lower_node_count", config.max_lower_node_count + ) + self.max_upper_node_count = _validate_sparse_grid_node_capacity( + "max_upper_node_count", config.max_upper_node_count + ) + strain_basis = config.strain_basis + collider_basis = config.collider_basis + strain_rebuild_safe = strain_basis[:3] == "pic" or strain_basis in ("P0", "P1d", "Q1d", "Q1") + collider_rebuild_safe = collider_basis[:3] == "pic" or collider_basis in ("Q1", "S2", "S3") + self._sparse_rebuildable = ( + self.grid_type == "sparse" + and self.max_active_cell_count > 0 + and self.grid_padding == 0 + and self.velocity_basis == "Q1" + and strain_rebuild_safe + and collider_rebuild_safe + ) + self._grid_status = None + self._grid_accumulated_status = None + self._grid_point_mask = None self.solver = _resolve_solver_spec(config.solver, self.velocity_basis) self.coloring = any("gauss-seidel" in solver or "gs" in solver for solver in self.solver) self.apic = config.transfer_scheme == "apic" self.gimp = config.integration_scheme == "gimp" - self.max_active_cell_count = config.max_active_cell_count - self.collider_normal_from_sdf_gradient = config.collider_normal_from_sdf_gradient self.collider_basis = config.collider_basis @@ -954,23 +1210,23 @@ def __init__( raise ValueError(f"Invalid collider velocity mode: {config.collider_velocity_mode}") self.collider_velocity_mode = config.collider_velocity_mode - if config.warmstart_mode == "none": - self._stress_warmstart = "" - elif config.warmstart_mode == "auto": - if self.strain_basis in ("P1d", "Q1d"): - self._stress_warmstart = "particles" - else: - self._stress_warmstart = "grid" - else: - if config.warmstart_mode not in ("particles", "grid", "smoothed"): - raise ValueError(f"Invalid warmstart mode: {config.warmstart_mode}") - self._stress_warmstart = config.warmstart_mode + warmstart_mode = config.warmstart_mode + if warmstart_mode not in ("none", "auto", "particles", "grid", "smoothed"): + raise ValueError(f"Invalid warmstart mode: {warmstart_mode}") + if warmstart_mode == "auto": + warmstart_mode = "particles" if self._sparse_rebuildable or self.strain_basis in ("P1d", "Q1d") else "grid" + if self._sparse_rebuildable and warmstart_mode in ("grid", "smoothed"): + raise ValueError( + f"Config.warmstart_mode={config.warmstart_mode!r} is not supported with rebuildable sparse grids " + "because their topology changes in place; use 'none', 'auto', or 'particles'." + ) + self._stress_warmstart = "" if warmstart_mode == "none" else warmstart_mode self._use_cuda_graph = self.model.device.is_cuda and wp.is_conditional_graph_supported() self._timers_use_nvtx = False - # Pre-allocate scratchpad and last step data so that step() can be graph-captured + # Materialize graph-persistent topology and buffers before callers can capture step(). self._scratchpad = None self._last_step_data = LastStepData() with wp.ScopedDevice(model.device): @@ -1001,6 +1257,7 @@ def setup_collider( body_mass: wp.array | None = None, body_inv_inertia: wp.array | None = None, body_q: wp.array | None = None, + collider_world_ids: list[int] | None = None, ) -> None: """Configure collider geometry and material properties. @@ -1018,9 +1275,25 @@ def setup_collider( collider_particle_ids: For deformable mesh colliders, model particle ids corresponding to each mesh vertex. model: The model to read collider properties from. Default to solver's model. body_com: For dynamic colliders, per-body center of mass. - body_mass: For dynamic colliders, per-body mass. Pass zeros for kinematic bodies. + body_mass: For dynamic colliders, per-body effective mass. When omitted, bodies flagged + with :attr:`newton.BodyFlags.KINEMATIC` have zero effective mass. An explicit array + is authoritative. body_inv_inertia: For dynamic colliders, per-body inverse inertia. body_q: For dynamic colliders, per-body initial transform. + collider_world_ids: Per-collider Newton world IDs. Custom meshes default to global + (``-1``). In isolated mode, body-backed colliders infer their body's world and + require any supplied ID to match. Shared-mode default discovery globalizes + colliders. IDs must be ``-1`` or in ``[0, model.world_count)``. + + .. experimental:: + + Per-world MPM collider filtering may change without prior notice. + + Raises: + ValueError: If collider-aligned inputs have different lengths, a world ID is invalid, + an isolated external collider model has a different ``world_count``, or an isolated + global body-backed collider is dynamic. Replicate a global dynamic collider per + world, make it static or kinematic, or disable ``Config.separate_worlds``. """ self._mpm_model.setup_collider( collider_meshes=collider_meshes, @@ -1035,6 +1308,7 @@ def setup_collider( body_mass=body_mass, body_inv_inertia=body_inv_inertia, body_q=body_q, + collider_world_ids=collider_world_ids, ) self._last_step_data.save_collider_current_position(self._mpm_model.collider_body_q) @@ -1044,6 +1318,304 @@ def voxel_size(self) -> float: """Grid voxel size used by the solver.""" return self._mpm_model.voxel_size + def _check_sparse_grid_rebuild_status(self) -> None: + """Raise if a rebuildable sparse grid exceeded its reserved capacity. + + The check synchronizes the solver device and is therefore intended for + initialization diagnostics or calls made after graph replay, not + from inside graph capture. + """ + if self._grid_status is None: + return + if self.model.device.is_capturing: + raise RuntimeError("Cannot inspect sparse grid rebuild status during graph capture") + + status = int(self._grid_status.numpy()[0]) + if self._grid_accumulated_status is not None: + status |= int(self._grid_accumulated_status.numpy()[0]) + if status != wp.Volume.REBUILD_SUCCESS: + raise _sparse_grid_rebuild_error(status) + + def check_status(self) -> None: + """Raise if a sparse-grid rebuild exceeded its reserved capacity. + + Rebuildable sparse-grid failures accumulate until a valid + :meth:`reset`. Call this method outside graph capture after replay; the + check synchronizes the solver device. It is a no-op when the solver has + no asynchronous sparse-grid status buffer. + + Raises: + RuntimeError: If rebuild status is inspected during graph capture, + or if a rebuild reported a capacity or topology failure. + """ + self._check_sparse_grid_rebuild_status() + + def _clear_sparse_grid_rebuild_status(self) -> None: + """Clear the latest and accumulated sparse-grid rebuild status. + + Call this outside graph capture after handling a reported rebuild + failure. A subsequent :meth:`check_status` then reports only failures + produced after this boundary. + + Raises: + RuntimeError: If called while the solver device is capturing a graph. + """ + if self.model.device.is_capturing: + raise RuntimeError("Cannot clear sparse grid rebuild status during graph capture") + if self._grid_status is not None: + self._grid_status.zero_() + if self._grid_accumulated_status is not None: + self._grid_accumulated_status.zero_() + + @staticmethod + def _validate_reset_array( + array: wp.array | None, + *, + name: str, + shape: tuple[int, ...], + dtype, + device: wp.Device, + ) -> None: + if not isinstance(array, wp.array): + raise ValueError(f"state.{name} must be a Warp array.") + if array.shape != shape: + raise ValueError(f"state.{name} has shape {array.shape}, expected {shape}.") + if array.dtype != dtype: + raise TypeError(f"state.{name} has dtype {array.dtype}, expected {dtype}.") + if array.device != device: + raise ValueError(f"state.{name} is on device {array.device}, expected {device}.") + + def _validate_reset_inputs(self, state: newton.State, world_mask: wp.array | None) -> None: + if state is None: + raise ValueError("'state' argument is required.") + + particle_count = int(self.model.particle_count) + if particle_count != self._initial_particle_count: + raise RuntimeError( + "SolverImplicitMPM model.particle_count changed after construction: " + f"expected {self._initial_particle_count}, got {particle_count}." + ) + world_count = int(self.model.world_count) + if world_count != self._initial_world_count: + raise RuntimeError( + "SolverImplicitMPM model.world_count changed after construction: " + f"expected {self._initial_world_count}, got {world_count}." + ) + + device = self.model.device + particle_shape = (self._initial_particle_count,) + self._validate_reset_array( + state.particle_q, + name="particle_q", + shape=particle_shape, + dtype=wp.vec3, + device=device, + ) + self._validate_reset_array( + state.particle_qd, + name="particle_qd", + shape=particle_shape, + dtype=wp.vec3, + device=device, + ) + if not hasattr(state, "mpm"): + raise ValueError("state is missing the 'mpm' custom-attribute namespace.") + for name, dtype in ( + ("particle_elastic_strain", wp.mat33), + ("particle_transform", wp.mat33), + ("particle_qd_grad", wp.mat33), + ("particle_stress", wp.mat33), + ("particle_Jp", wp.float32), + ): + self._validate_reset_array( + getattr(state.mpm, name, None), + name=f"mpm.{name}", + shape=particle_shape, + dtype=dtype, + device=device, + ) + + if self.model.body_count > 0: + self._validate_reset_array( + state.body_q, + name="body_q", + shape=(self.model.body_count,), + dtype=wp.transform, + device=device, + ) + + if world_mask is None: + return + if not isinstance(world_mask, wp.array): + raise TypeError("world_mask must be a Warp array with dtype wp.bool.") + expected_shape = (self._initial_world_count + 1,) + if world_mask.shape != expected_shape: + raise ValueError(f"world_mask has shape {world_mask.shape}, expected {expected_shape}.") + if world_mask.dtype != wp.bool: + raise TypeError(f"world_mask has dtype {world_mask.dtype}, expected {wp.bool}.") + if world_mask.device != device: + raise ValueError(f"world_mask is on device {world_mask.device}, expected {device}.") + + def _reset_grid_warmstart_partition(self, name, field, scratch_field) -> fem.SpacePartition: + scratch_partition = scratch_field.space_partition + if field.space.topology == scratch_partition.space_topology: + partition = scratch_partition + else: + partition = fem.make_space_partition( + space_topology=field.space.topology, + geometry_partition=scratch_partition.geo_partition, + with_halo=False, + max_node_count=field.dof_values.shape[0], + environment_first=True, + device=self.model.device, + temporary_store=self.temporary_store, + ) + + environment_offsets = getattr(partition, "env_offsets", None) + expected_shape = (self._initial_world_count + 1,) + if environment_offsets is None or environment_offsets.shape != expected_shape: + raise RuntimeError( + f"Masked reset cannot selectively clear last-step {name}: its grid partition does not expose " + f"environment offsets with shape {expected_shape}." + ) + return partition + + def _validate_reset_warmstart_fields( + self, world_mask: wp.array | None + ) -> tuple[fem.SpacePartition | None, fem.SpacePartition | None]: + reset_partitions = [] + for name, field, scratch_field in ( + ("ws_impulse_field", self._last_step_data.ws_impulse_field, self._scratchpad.impulse_field), + ("ws_stress_field", self._last_step_data.ws_stress_field, self._scratchpad.stress_field), + ): + if field is None: + reset_partitions.append(None) + continue + + expected_shape = (field.space_partition.node_count(),) + if field.dof_values.shape != expected_shape: + raise ValueError( + f"last-step {name}.dof_values has shape {field.dof_values.shape}, expected {expected_shape}." + ) + + if world_mask is None or isinstance(field.space.basis, fem.PointBasisSpace): + reset_partitions.append(None) + continue + if self._initial_world_count > 1 and not self._separate_worlds: + raise RuntimeError( + "Masked reset cannot selectively clear grid-backed warm starts when " + "Config.separate_worlds=False for a multi-world model; set separate_worlds=True or reset all worlds." + ) + reset_partitions.append(self._reset_grid_warmstart_partition(name, field, scratch_field)) + + return tuple(reset_partitions) + + def _clear_reset_warmstarts( + self, + world_mask: wp.array | None, + reset_partitions: tuple[fem.SpacePartition | None, fem.SpacePartition | None], + ) -> None: + for field, partition in zip( + (self._last_step_data.ws_impulse_field, self._last_step_data.ws_stress_field), + reset_partitions, + strict=True, + ): + if field is None: + continue + if world_mask is None: + field.dof_values.zero_() + elif self._initial_particle_count > 0 and isinstance(field.space.basis, fem.PointBasisSpace): + wp.launch( + reset_mpm_point_warmstart, + dim=self._initial_particle_count, + inputs=[self.model.particle_world, world_mask, field.dof_values], + device=self.model.device, + ) + elif partition is not None and partition.node_count() > 0: + wp.launch( + reset_mpm_grid_warmstart, + dim=partition.node_count(), + inputs=[world_mask, partition.env_offsets, partition.space_node_indices(), field.dof_values], + device=self.model.device, + ) + + @override + def reset( + self, + state: newton.State, + world_mask: wp.array | None = None, + flags: StateFlags | int | None = None, + ) -> None: + """Reset implicit MPM history for all or selected worlds. + + Particle history is reset when ``flags`` includes either + :attr:`~newton.StateFlags.PARTICLE_Q` or + :attr:`~newton.StateFlags.PARTICLE_QD`. If ``flags`` is ``None``, all + particle history is reset. Particle- and grid-backed warm starts are + cleared for selected worlds. Grid-backed warm starts cannot be + selectively cleared on a shared multi-world grid; use + :attr:`Config.separate_worlds` or a full reset. A full reset clears + every warm-start field. Sparse-grid rebuild status is always cleared at + a valid reset boundary, and the previous-collider-pose cache is + refreshed from ``state``. The final mask entry selects global + particle-backed history and collider poses whose world index is ``-1``. + + Args: + state: Simulation state whose MPM history is modified in place. + world_mask: Optional boolean mask of shape + ``(model.world_count + 1,)`` selecting worlds to reset. Entries + before the last select local worlds by index, and the last + entry selects global objects whose world index is ``-1``. If + ``None``, reset all worlds and global objects. + + .. experimental:: + + Selective per-world MPM reset behavior may change without prior notice. + flags: Optional state bitmask. If ``None``, reset all particle + history. + """ + self._validate_reset_inputs(state, world_mask) + reset_partitions = self._validate_reset_warmstart_fields(world_mask) + state_flags = int(StateFlags.ALL if flags is None else flags) + reset_particle_history = bool(state_flags & int(StateFlags.PARTICLE)) + + # Clearing first ensures a capture-time rejection cannot leave state + # partially reset. + self._clear_sparse_grid_rebuild_status() + + with wp.ScopedDevice(self.model.device): + self._clear_reset_warmstarts(world_mask, reset_partitions) + + if reset_particle_history and self.model.particle_count > 0: + if world_mask is None: + identity = wp.mat33(np.eye(3)) + state.mpm.particle_elastic_strain.fill_(identity) + state.mpm.particle_transform.fill_(identity) + state.mpm.particle_qd_grad.zero_() + state.mpm.particle_stress.zero_() + state.mpm.particle_Jp.fill_(1.0) + else: + wp.launch( + reset_mpm_particle_history, + dim=self.model.particle_count, + inputs=[ + self.model.particle_world, + world_mask, + state.mpm.particle_elastic_strain, + state.mpm.particle_transform, + state.mpm.particle_qd_grad, + state.mpm.particle_stress, + state.mpm.particle_Jp, + ], + device=self.model.device, + ) + + self._last_step_data.save_collider_current_position( + state.body_q, + body_world=self.model.body_world, + world_mask=world_mask, + ) + @override def step( self, @@ -1373,6 +1945,7 @@ def project_outside(self, state_in: newton.State, state_out: newton.State, dt: f state_in.mpm.particle_qd_grad, self._mpm_model.particle_flags, self.model.particle_mass, + self._particle_environment, self._mpm_model.collider, state_in.body_q, state_in.body_qd if self.collider_velocity_mode == "forward" else None, @@ -1450,7 +2023,16 @@ def update_render_grains( dt: Time step duration. """ - return update_render_grains(state_prev, state, grains, self._mpm_model.particle_radius, dt) + with wp.ScopedDevice(grains.device): + return update_render_grains( + state_prev, + state, + grains, + self._mpm_model.particle_radius, + dt, + particle_environment=self._particle_environment, + temporary_store=self.temporary_store, + ) def _allocate_grid( self, @@ -1478,8 +2060,94 @@ def _allocate_grid( """ with self._timer("Allocate grid"): if self.grid_type == "sparse": - volume = allocate_by_voxels(positions, voxel_size, padding_voxels=padding_voxels) - grid = fem.Nanogrid(volume, temporary_store=temporary_store) + if self._separate_worlds: + if self._sparse_rebuildable: + if self._grid_status is None: + self._grid_status = wp.zeros(1, dtype=wp.uint32, device=positions.device) + self._grid_accumulated_status = wp.zeros(1, dtype=wp.uint32, device=positions.device) + point_mask = self._update_grid_point_mask(positions, self._mpm_model.particle_flags) + guard_cells = 3 + capacity_kwargs = _rebuild_capacity( + positions, + voxel_size, + 16.0, + self.max_active_cell_count, + point_mask=point_mask, + point_environment=self._particle_environment, + environment_count=self._environment_count, + guard_cells=guard_cells, + temporary_store=temporary_store, + max_leaf_node_count=self.max_leaf_node_count, + max_lower_node_count=self.max_lower_node_count, + max_upper_node_count=self.max_upper_node_count, + ) + grid = fem.Nanogrid.from_environment_voxels( + positions, + self._particle_environment, + self._environment_count, + point_mask=point_mask, + voxel_size=voxel_size, + temporary_store=temporary_store, + device=positions.device, + rebuildable=True, + status=self._grid_status, + **capacity_kwargs, + ) + self._check_sparse_grid_rebuild_status() + else: + cell_ijks = [ + voxel_coordinates(positions[begin:end], voxel_size, padding_voxels=padding_voxels) + if begin != end + else wp.empty(0, dtype=wp.vec3i, device=positions.device) + for begin, end in self._particle_world_ranges + ] + cell_count = sum(cell_ijk.shape[0] for cell_ijk in cell_ijks) + cell_ijk = wp.empty(cell_count, dtype=wp.vec3i, device=positions.device) + cell_environment = wp.empty(cell_count, dtype=wp.int32, device=positions.device) + cell_offset = 0 + for environment, environment_cell_ijk in enumerate(cell_ijks): + environment_cell_count = environment_cell_ijk.shape[0] + if environment_cell_count: + wp.copy( + cell_ijk, + environment_cell_ijk, + dest_offset=cell_offset, + count=environment_cell_count, + ) + cell_environment[cell_offset : cell_offset + environment_cell_count].fill_(environment) + cell_offset += environment_cell_count + grid = fem.Nanogrid.from_environment_voxels( + cell_ijk, + cell_environment, + self._environment_count, + voxel_size=voxel_size, + temporary_store=temporary_store, + device=positions.device, + ) + else: + point_mask = None + if self._sparse_rebuildable: + if self._grid_status is None: + self._grid_status = wp.zeros(1, dtype=wp.uint32, device=positions.device) + self._grid_accumulated_status = wp.zeros(1, dtype=wp.uint32, device=positions.device) + point_mask = self._update_grid_point_mask(positions, self._mpm_model.particle_flags) + volume = allocate_by_voxels( + positions, + voxel_size, + padding_voxels=padding_voxels, + rebuildable=self._sparse_rebuildable, + max_active_voxels=self.max_active_cell_count if self._sparse_rebuildable else None, + status=self._grid_status, + point_mask=point_mask, + max_leaf_node_count=self.max_leaf_node_count, + max_lower_node_count=self.max_lower_node_count, + max_upper_node_count=self.max_upper_node_count, + ) + if self._sparse_rebuildable: + self._check_sparse_grid_rebuild_status() + grid = fem.Nanogrid(volume, temporary_store=temporary_store, rebuildable=True) + else: + grid = fem.Nanogrid(volume, temporary_store=temporary_store) else: # Compute bounds and transfer to host device = positions.device @@ -1520,10 +2188,33 @@ def _allocate_grid( bounds_lo=wp.vec3(grid_min * voxel_size), bounds_hi=wp.vec3(grid_max * voxel_size), res=wp.vec3i((grid_max - grid_min).astype(int)), + env_count=self._environment_count, ) return grid + def _update_grid_point_mask( + self, + positions: wp.array[wp.vec3], + particle_flags: wp.array[wp.int32], + ) -> wp.array[wp.int32]: + if self._grid_point_mask is None: + self._grid_point_mask = wp.empty( + shape=particle_flags.shape, + dtype=wp.int32, + device=particle_flags.device, + ) + elif self._grid_point_mask.shape != particle_flags.shape: + raise RuntimeError("Implicit MPM particle count changed after sparse grid initialization") + + wp.launch( + build_active_particle_mask, + dim=particle_flags.shape[0], + inputs=[positions, particle_flags, self._grid_point_mask], + device=particle_flags.device, + ) + return self._grid_point_mask + def _create_geometry_partition( self, grid: fem.Geometry, positions: wp.array, particle_flags: wp.array, max_cell_count: int ): @@ -1531,15 +2222,26 @@ def _create_geometry_partition( active_cells = fem.borrow_temporary(self.temporary_store, shape=grid.cell_count(), dtype=int) active_cells.zero_() - fem.interpolate( - mark_active_cells, - dim=positions.shape[0], - at=fem.Cells(grid), - values={ + if self._separate_worlds: + active_cell_integrand = mark_active_cells_by_environment + active_cell_values = { "positions": positions, "particle_flags": particle_flags, + "particle_environment": self._particle_environment, "active_cells": active_cells, - }, + } + else: + active_cell_integrand = mark_active_cells + active_cell_values = { + "positions": positions, + "particle_flags": particle_flags, + "active_cells": active_cells, + } + fem.interpolate( + active_cell_integrand, + dim=positions.shape[0], + at=fem.Cells(grid), + values=active_cell_values, temporary_store=self.temporary_store, ) @@ -1574,6 +2276,7 @@ def _rebuild_scratchpad(self, pic: fem.PicQuadrature): velocity_basis_str=self.velocity_basis, collider_basis_str=self.collider_basis, max_cell_count=self.max_active_cell_count, + environment_first=self._separate_worlds, temporary_store=self.temporary_store, ) @@ -1595,8 +2298,25 @@ def _particles_to_cells(self, positions: wp.array) -> fem.PicQuadrature: # Rebuild grid - if self._scratchpad is not None and self.grid_type == "fixed": + # The fixed grid and the rebuildable sparse grid both persist across steps: the + # fixed grid is static, the sparse grid is refreshed in place from the current + # particles. Plain sparse (no rebuild support) reallocates the grid each step. + if self._scratchpad is not None and (self.grid_type == "fixed" or self._sparse_rebuildable): grid = self._scratchpad.grid + if self._sparse_rebuildable: + point_mask = self._update_grid_point_mask(positions, self._mpm_model.particle_flags) + grid.rebuild( + positions, + point_envs=self._particle_environment, + status=self._grid_status, + point_mask=point_mask, + ) + wp.launch( + record_volume_rebuild_status, + dim=1, + inputs=[self._grid_status, self._grid_accumulated_status], + device=positions.device, + ) else: grid = self._allocate_grid( positions, @@ -1606,9 +2326,11 @@ def _particles_to_cells(self, positions: wp.array) -> fem.PicQuadrature: padding_voxels=self.grid_padding, ) - # Build active partition + # Build active partition. Plain sparse uses the whole grid; fixed and rebuildable + # sparse use a capacity-bounded partition that masks to the active cells (the + # rebuildable grid's cell buffers are capacity-sized and include unused slots). with self._timer("Build active partition"): - if self.grid_type == "sparse": + if self.grid_type == "sparse" and not self._sparse_rebuildable: max_cell_count = -1 geo_partition = grid else: @@ -1621,65 +2343,42 @@ def _particles_to_cells(self, positions: wp.array) -> fem.PicQuadrature: with self._timer("Bin particles"): domain = fem.Cells(geo_partition) + # Whole-grid sparse domains use identical geometry and domain indices. Keeping + # them geometry-scoped also gives their PicQuadrature a distinct cache name + # from fixed/rebuildable grids, whose explicit partitions require domain indices. + use_domain_element_indices = not ( + self._separate_worlds and self.grid_type == "sparse" and not self._sparse_rebuildable + ) + if self.gimp: particle_locations = self._particle_grid_locations_gimp( - domain, positions, self._mpm_model.particle_radius + domain, positions, self._mpm_model.particle_radius, self._particle_environment + ) + pic = fem.PicQuadrature( + domain=domain, + positions=particle_locations, + measures=self._mpm_model.particle_volume, + temporary_store=self.temporary_store, + use_domain_element_indices=use_domain_element_indices, ) else: - particle_locations = self._particle_grid_locations(domain, positions) - - pic = fem.PicQuadrature( - domain=domain, - positions=particle_locations, - measures=self._mpm_model.particle_volume, - temporary_store=self.temporary_store, - use_domain_element_indices=True, - ) + pic = fem.PicQuadrature( + domain=domain, + positions=positions, + env_indices=self._particle_environment, + measures=self._mpm_model.particle_volume, + temporary_store=self.temporary_store, + use_domain_element_indices=use_domain_element_indices, + ) return pic - def _particle_grid_locations(self, domain: fem.GeometryDomain, positions: wp.array) -> wp.array: - """Convert particle positions to grid locations.""" - - cell_lookup = domain.element_partition_lookup - - @fem.cache.dynamic_kernel(suffix=domain.name) - def particle_locations( - cell_arg_value: domain.ElementArg, - domain_index_arg_value: domain.ElementIndexArg, - positions: wp.array[wp.vec3], - cell_index: wp.array[fem.ElementIndex], - cell_coords: wp.array[fem.Coords], - ): - p = wp.tid() - domain_arg = domain.DomainArg(cell_arg_value, domain_index_arg_value) - - sample = cell_lookup(domain_arg, positions[p]) - - cell_index[p] = domain.element_partition_index(domain_index_arg_value, sample.element_index) - cell_coords[p] = sample.element_coords - - device = positions.device - - cell_indices = fem.borrow_temporary(self.temporary_store, shape=positions.shape[0], dtype=fem.ElementIndex) - cell_coords = fem.borrow_temporary(self.temporary_store, shape=positions.shape[0], dtype=fem.Coords) - wp.launch( - particle_locations, - dim=positions.shape[0], - inputs=[ - domain.element_arg_value(device=device), - domain.element_index_arg_value(device=device), - positions, - cell_indices, - cell_coords, - ], - device=device, - ) - - return cell_indices, cell_coords - def _particle_grid_locations_gimp( - self, domain: fem.GeometryDomain, positions: wp.array, radii: wp.array + self, + domain: fem.GeometryDomain, + positions: wp.array, + radii: wp.array, + particle_environment: wp.array | None, ) -> wp.array: """Convert particle positions to grid locations.""" @@ -1706,12 +2405,15 @@ def add_cell( particle_cell_fractions[i] += cell_weight return - @fem.cache.dynamic_kernel(suffix=domain.name) + separate_worlds = self._separate_worlds + + @fem.cache.dynamic_kernel(suffix=f"{domain.name}_{'isolated' if separate_worlds else 'shared'}") def particle_locations_gimp( cell_arg_value: domain.ElementArg, domain_index_arg_value: domain.ElementIndexArg, positions: wp.array[wp.vec3], radii: wp.array[float], + particle_environment: wp.array[int], cell_index: wp.array2d[fem.ElementIndex], cell_coords: wp.array2d[fem.Coords], cell_fractions: wp.array2d[float], @@ -1732,7 +2434,10 @@ def particle_locations_gimp( k = vtx & 1 pos = center - wp.vec3(radius) + 2.0 * radius * wp.vec3(float(i), float(j), float(k)) - sample = cell_lookup(domain_arg, pos) + if wp.static(separate_worlds): + sample = cell_lookup(domain_arg, pos, int(particle_environment[p])) + else: + sample = cell_lookup(domain_arg, pos) if sample.element_index == fem.NULL_ELEMENT_INDEX: continue @@ -1773,6 +2478,7 @@ def particle_locations_gimp( domain.element_index_arg_value(device=device), positions, radii, + particle_environment, cell_indices, cell_coords, cell_fractions, @@ -1963,6 +2669,7 @@ def _rasterize_colliders( scratch.collider_adhesion, scratch.collider_ids, temporary_store=self.temporary_store, + node_environment_offsets=scratch.collider_environment_offsets, ) # normal interpolation @@ -2411,9 +3118,10 @@ def _solve_rheology( elastic_strain_delta=scratch.elastic_strain_delta_field.dof_values, plastic_strain_delta=scratch.plastic_strain_delta_field.dof_values, stress=scratch.stress_field.dof_values, + strain_environment_offsets=scratch.strain_environment_offsets, has_viscosity=self._mpm_model.has_viscosity, has_dilatancy=self._mpm_model.has_dilatancy, - strain_velocity_node_count=self._velocity_nodes_per_strain_sample, + strain_velocity_node_count=-1 if self._separate_worlds else self._velocity_nodes_per_strain_sample, ) collision_data = CollisionData( collider_mat=scratch.collider_matrix, @@ -2608,10 +3316,16 @@ def _warmstart_fields( domain = scratch.velocity_test.domain + # The rebuildable sparse grid is refreshed in place, so the previous step's grid + # topology no longer exists: a grid-to-grid (nonconforming) warmstart would read + # stale cells. Skip those transfers (they only accelerate convergence); the + # particle/point paths below still apply since they go through the PIC quadrature. + grid_to_grid_warmstart = not self._sparse_rebuildable + if isinstance(prev_impulse_field.space.basis, fem.PointBasisSpace): # point-based collisions, simply copy the previous impulses scratch.impulse_field.dof_values.assign(prev_impulse_field.dof_values[pic.cell_particle_indices]) - else: + elif grid_to_grid_warmstart: # Interpolate previous impulse prev_impulse_field = fem.NonconformingField( domain, prev_impulse_field, background=scratch.background_impulse_field @@ -2623,11 +3337,13 @@ def _warmstart_fields( reduction="first", temporary_store=self.temporary_store, ) + else: + scratch.impulse_field.dof_values.zero_() # Interpolate previous stress if isinstance(prev_stress_field.space.basis, fem.PointBasisSpace): scratch.stress_field.dof_values.assign(prev_stress_field.dof_values[pic.cell_particle_indices]) - elif self._stress_warmstart in ("grid", "smoothed"): + elif self._stress_warmstart in ("grid", "smoothed") and grid_to_grid_warmstart: prev_stress_field = fem.NonconformingField( domain, prev_stress_field, background=scratch.background_stress_field ) @@ -2638,6 +3354,11 @@ def _warmstart_fields( reduction="first", temporary_store=self.temporary_store, ) + elif not self._sparse_rebuildable: + pass + else: + # No grid-to-grid stress warmstart available for the rebuilt grid; start cold. + scratch.stress_field.dof_values.zero_() def _save_for_next_warmstart( self, scratch: ImplicitMPMScratchpad, pic: fem.PicQuadrature, last_step_data: LastStepData diff --git a/newton/_src/solvers/kamino/_src/__init__.py b/newton/_src/solvers/kamino/_src/__init__.py index 00c52235e1..60af581a4e 100644 --- a/newton/_src/solvers/kamino/_src/__init__.py +++ b/newton/_src/solvers/kamino/_src/__init__.py @@ -12,8 +12,13 @@ convert_geom_offset_origin_to_com, ) from .core.control import ControlKamino -from .core.conversions import convert_model_joint_transforms -from .core.gravity import convert_model_gravity +from .core.conversions import ( + compute_material_first_shape, + convert_model_joint_actuation, + convert_model_joint_transforms, + convert_model_materials, + validate_model_joint_updates, +) from .core.joints import JOINT_QMAX, JOINT_QMIN, JointActuationType from .core.model import ModelKamino from .core.state import StateKamino @@ -40,13 +45,16 @@ "ModelKamino", "SolverKaminoImpl", "StateKamino", + "compute_material_first_shape", "convert_base_origin_to_com", "convert_body_com_to_origin", "convert_body_origin_to_com", "convert_contacts_kamino_to_newton", "convert_contacts_newton_to_kamino", "convert_geom_offset_origin_to_com", - "convert_model_gravity", + "convert_model_joint_actuation", "convert_model_joint_transforms", + "convert_model_materials", "msg", + "validate_model_joint_updates", ] diff --git a/newton/_src/solvers/kamino/_src/core/builder.py b/newton/_src/solvers/kamino/_src/core/builder.py index fd77e19540..8b9ac197dc 100644 --- a/newton/_src/solvers/kamino/_src/core/builder.py +++ b/newton/_src/solvers/kamino/_src/core/builder.py @@ -31,7 +31,7 @@ from .shapes import ShapeDescriptorType, max_contacts_for_shape_pair from .size import SizeKamino from .time import TimeModel -from .types import to_warp_int32_array +from .types import ArrayLike, to_warp_int32_array from .world import WorldDescriptor ### @@ -206,7 +206,7 @@ def up_axes(self) -> list[Axis]: @property def gravity(self) -> list[GravityDescriptor]: - """Returns the list of gravity descriptors for each world contained in the model.""" + """Returns the gravity descriptor for each world contained in the model.""" return self._gravity @property @@ -261,7 +261,7 @@ def add_world( name: str = "world", uid: str | None = None, up_axis: Axis | None = None, - gravity: GravityDescriptor | None = None, + gravity: GravityDescriptor | ArrayLike | None = None, ) -> int: """ Add a new world to the model. @@ -272,8 +272,8 @@ def add_world( If None, a UUID will be generated. up_axis: The up axis of the world. If None, Axis.Z will be used. - gravity: The gravity descriptor of the world. - If None, a default gravity descriptor will be used. + gravity: The gravity descriptor or vector [m/s²] of the world. + If ``None``, Newton's default gravity is used along the negative up axis. Returns: The index of the newly added world. @@ -293,7 +293,9 @@ def add_world( # Set gravity if gravity is None: - gravity = GravityDescriptor() + gravity = GravityDescriptor.default_from_up_axis(up_axis) + elif not isinstance(gravity, GravityDescriptor): + gravity = GravityDescriptor.from_array(gravity) self._gravity.append(gravity) # Register the default material in the new world @@ -752,26 +754,20 @@ def set_up_axis(self, axis: Axis, world_index: int = 0): # Set the new up axis self._up_axes[world_index] = axis - def set_gravity(self, gravity: GravityDescriptor, world_index: int = 0): + def set_gravity(self, gravity: GravityDescriptor | ArrayLike, world_index: int = 0): """ - Set the gravity descriptor for a specific world. + Set the gravity vector for a specific world. Args: - gravity: The new gravity descriptor to be set. - world_index: The index of the world for which to set the gravity descriptor. + gravity: The new gravity descriptor or vector [m/s²]. + world_index: The index of the world for which to set gravity. Defaults to the first world with index `0`. - - Raises: - TypeError: If the provided gravity descriptor is not of type `GravityDescriptor`. """ # Check if the world index is valid self._check_world_index(world_index) - # Check if the gravity descriptor is valid if not isinstance(gravity, GravityDescriptor): - raise TypeError(f"Invalid gravity descriptor type: {type(gravity)}. Must be `GravityDescriptor`.") - - # Set the new gravity configurations + gravity = GravityDescriptor.from_array(gravity) self._gravity[world_index] = gravity def set_default_material(self, material: MaterialDescriptor, world_index: int = 0): @@ -981,7 +977,6 @@ def finalize( info_base_jid = [] # Initialize the gravity data collections - gravity_g_dir_acc = [] gravity_vector = [] # Initialize the body data collections @@ -1101,8 +1096,7 @@ def collect_model_info_data(): # A helper function to collect model gravity data def collect_gravity_model_data(): for w in range(num_worlds): - gravity_g_dir_acc.append(self._gravity[w].dir_accel()) - gravity_vector.append(self._gravity[w].vector()) + gravity_vector.append(self._gravity[w].vector) # A helper function to collect model bodies data def collect_body_model_data(): @@ -1364,10 +1358,7 @@ def collect_material_pairs_model_data(): ) # Construct model gravity data - model_gravity = GravityModel( - g_dir_acc=wp.array(gravity_g_dir_acc, dtype=wp.vec4f), - vector=wp.array(gravity_vector, dtype=wp.vec4f, requires_grad=requires_grad), - ) + model_gravity = GravityModel(vector=wp.array(gravity_vector, dtype=wp.vec3, requires_grad=requires_grad)) # Create the bodies model model_bodies = RigidBodiesModel( @@ -1766,7 +1757,7 @@ def compute_required_contact_capacity( else: world_max_contacts[geom1.wid] += num_contacts - # Override the per-world maximum contacts if specified in the settings + # Cap per-world totals when a per-world maximum is specified if max_contacts_per_world is not None: for w in range(self.num_worlds): world_max_contacts[w] = min(world_max_contacts[w], max_contacts_per_world) diff --git a/newton/_src/solvers/kamino/_src/core/conversions.py b/newton/_src/solvers/kamino/_src/core/conversions.py index e75df9ff9b..1ee8dcde23 100644 --- a/newton/_src/solvers/kamino/_src/core/conversions.py +++ b/newton/_src/solvers/kamino/_src/core/conversions.py @@ -18,9 +18,11 @@ convert_body_origin_to_com, convert_geom_offset_origin_to_com, ) -from .builder import JointActuationType from .geometry import GeometriesModel from .joints import ( + JOINT_QMAX, + JOINT_QMIN, + JointActuationType, JointDoFType, JointsModel, ) @@ -39,10 +41,13 @@ __all__ = [ "convert_geometries", "convert_joints", + "convert_model_joint_actuation", "convert_model_joint_transforms", + "convert_model_materials", "convert_rigid_bodies", "convert_target_coords_to_target_dofs", "convert_target_dofs_to_target_coords", + "validate_model_joint_updates", ] @@ -57,6 +62,35 @@ ### +@wp.func +def joint_actuation_type_from_dofs( + dof_start: int, + dof_end: int, + target_mode: wp.array[wp.int32], +) -> int: + """Aggregate Newton's per-DoF target modes into a Kamino joint actuation type.""" + joint_target_mode = int(0) + for dof in range(dof_start, dof_end): + joint_target_mode = max(joint_target_mode, target_mode[dof]) + return JointActuationType.from_newton_wp(joint_target_mode) + + +@wp.func +def joint_requires_dynamic_constraints( + dof_start: int, + dof_end: int, + armature: wp.array[wp.float32], + damping: wp.array[wp.float32], + target_ke: wp.array[wp.float32], + target_kd: wp.array[wp.float32], +) -> bool: + """Return whether any DoF makes a joint dynamic.""" + dynamic = bool(False) + for dof in range(dof_start, dof_end): + dynamic = dynamic or (armature[dof] > 0.0 or damping[dof] > 0.0 or target_ke[dof] > 0.0 or target_kd[dof] > 0.0) + return dynamic + + @wp.kernel def world_max_contacts_kernel( # Inputs: @@ -96,6 +130,150 @@ def world_max_contacts_kernel( wp.atomic_add(world_max_contacts, world_id, num_contacts) +@wp.kernel +def material_first_shape_kernel( + # Inputs: + geom_material: wp.array[wp.int32], + # Outputs: + first_shape: wp.array[wp.int32], +): + """Record the first shape index associated with each material.""" + shape = wp.tid() + material = geom_material[shape] + if material >= 0: + wp.atomic_min(first_shape, material, shape) + + +@wp.kernel +def validate_material_update_kernel( + shape_friction: wp.array[wp.float32], + shape_restitution: wp.array[wp.float32], + geom_material: wp.array[wp.int32], + first_shape: wp.array[wp.int32], + conflict_material: wp.array[wp.int32], +): + """Find the first material whose shapes have conflicting properties.""" + shape = wp.tid() + material = geom_material[shape] + if material < 0: + return + representative = first_shape[material] + if ( + shape_friction[shape] != shape_friction[representative] + or shape_restitution[shape] != shape_restitution[representative] + ): + wp.atomic_min(conflict_material, 0, material) + + +@wp.kernel +def update_materials_kernel( + # Inputs: + shape_friction: wp.array[wp.float32], + shape_restitution: wp.array[wp.float32], + first_shape: wp.array[wp.int32], + shape_count: int, + # Outputs: + restitution: wp.array[wp.float32], + static_friction: wp.array[wp.float32], + dynamic_friction: wp.array[wp.float32], + pair_restitution: wp.array[wp.float32], + pair_static_friction: wp.array[wp.float32], + pair_dynamic_friction: wp.array[wp.float32], +): + """Update Kamino material properties from cached representative shapes. + + The material-zero properties are also copied to the default material pair. + """ + material = wp.tid() + shape = first_shape[material] + if shape < shape_count: + friction = shape_friction[shape] + restitution[material] = shape_restitution[shape] + static_friction[material] = friction + dynamic_friction[material] = friction + if material == 0: + pair_restitution[0] = shape_restitution[shape] + pair_static_friction[0] = friction + pair_dynamic_friction[0] = friction + + +@wp.kernel +def validate_joint_dof_updates_kernel( + # Inputs: + joint_qd_start: wp.array[wp.int32], + joint_armature: wp.array[wp.float32], + joint_damping: wp.array[wp.float32], + joint_target_ke: wp.array[wp.float32], + joint_target_kd: wp.array[wp.float32], + num_dynamic_cts: wp.array[wp.int32], + joint_limit_lower: wp.array[wp.float32], + joint_limit_upper: wp.array[wp.float32], + built_limit_finite: wp.array[wp.int32], + joint_count: int, + dof_count: int, + # Outputs: + violations: wp.array[wp.int32], +): + """Find the first structural change to joint degree-of-freedom properties.""" + tid = wp.tid() + if tid < joint_count: + dof_start = joint_qd_start[tid] + dof_end = joint_qd_start[tid + 1] + if joint_requires_dynamic_constraints( + dof_start, + dof_end, + joint_armature, + joint_damping, + joint_target_ke, + joint_target_kd, + ) != (num_dynamic_cts[tid] > 0): + wp.atomic_min(violations, 0, tid) + + if tid < dof_count: + current_finite = joint_limit_lower[tid] > JOINT_QMIN or joint_limit_upper[tid] < JOINT_QMAX + if current_finite != (built_limit_finite[tid] != 0): + wp.atomic_min(violations, 1, tid) + + +@wp.kernel +def validate_joint_actuation_updates_kernel( + # Inputs: + joint_qd_start: wp.array[wp.int32], + joint_target_mode: wp.array[wp.int32], + act_type: wp.array[wp.int32], + # Outputs: + violations: wp.array[wp.int32], +): + """Find the first joint with an invalid or structurally changed actuation type.""" + joint = wp.tid() + current_actuation = joint_actuation_type_from_dofs( + joint_qd_start[joint], + joint_qd_start[joint + 1], + joint_target_mode, + ) + if current_actuation < 0: + wp.atomic_min(violations, 3, joint) + elif (current_actuation == JointActuationType.PASSIVE) != (act_type[joint] == JointActuationType.PASSIVE): + wp.atomic_min(violations, 2, joint) + + +@wp.kernel +def update_joint_actuation_kernel( + # Inputs: + joint_qd_start: wp.array[wp.int32], + joint_target_mode: wp.array[wp.int32], + # Outputs: + act_type: wp.array[wp.int32], +): + """Update each joint's Kamino actuation type from its target modes.""" + joint = wp.tid() + act_type[joint] = joint_actuation_type_from_dofs( + joint_qd_start[joint], + joint_qd_start[joint + 1], + joint_target_mode, + ) + + @wp.kernel def rigid_bodies_indexing_kernel( # Inputs: @@ -188,21 +366,19 @@ def joint_conversion_kernel( joint_num_dofs[joint_id] = ndofs_j # Determine Kamino actuation mode for joint - joint_dofs_target_mode_j = int(0) - for dof_id in range(ndofs_j): - joint_dofs_target_mode_j = max(joint_dofs_target_mode_j, model_joint_target_mode[dofs_start_j + dof_id]) - act_type_j = JointActuationType.from_newton_wp(joint_dofs_target_mode_j) + act_type_j = joint_actuation_type_from_dofs(dofs_start_j, dofs_start_j + ndofs_j, model_joint_target_mode) assert act_type_j >= 0, "Joint actuation type must be valid" joint_act_type[joint_id] = act_type_j - is_dynamic_j = bool(False) # Infer if the joint requires dynamic constraints - for dof_id in range(ndofs_j): - a_j = model_joint_armature[dofs_start_j + dof_id] - b_j = model_joint_damping[dofs_start_j + dof_id] - ke_j = model_joint_target_ke[dofs_start_j + dof_id] - kd_j = model_joint_target_kd[dofs_start_j + dof_id] - is_dynamic_j = is_dynamic_j or (a_j > 0.0) or (b_j > 0.0) or (ke_j > 0.0) or (kd_j > 0.0) + is_dynamic_j = joint_requires_dynamic_constraints( + dofs_start_j, + dofs_start_j + ndofs_j, + model_joint_armature, + model_joint_damping, + model_joint_target_ke, + model_joint_target_kd, + ) # Set joint dimensions joint_num_kinematic_cts[joint_id] = ncts_j @@ -579,8 +755,8 @@ def compute_required_contact_capacity( max_contacts_per_pair: Optional maximum number of contacts to allocate per shape pair. If `None`, no per-pair limit is applied. max_contacts_per_world: Optional maximum number of contacts to allocate per world. - If `None`, no per-world limit is applied, otherwise it will - override the computed per-world requirements if it is larger. + If `None`, no per-world limit is applied, otherwise caps the computed + per-world requirements at this value. Returns: (model_required_contacts, world_required_contacts): @@ -613,7 +789,7 @@ def compute_required_contact_capacity( ) world_max_contacts = world_max_contacts_wp.numpy() - # Override the per-world maximum contacts if specified in the settings + # Cap per-world totals when a per-world maximum is specified if max_contacts_per_world is not None: world_max_contacts = np.minimum(world_max_contacts, max_contacts_per_world) @@ -621,6 +797,98 @@ def compute_required_contact_capacity( return int(np.sum(world_max_contacts)), world_max_contacts.astype(int).tolist() +def validate_model_joint_updates( + model: Model, + joints: JointsModel, + built_limit_finite: wp.array[wp.int32], + violations: wp.array[wp.int32], + *, + check_dof: bool, + check_actuation: bool, +) -> int: + """Validate that runtime joint edits preserve Kamino's structural layout. + + ``violations`` is a four-entry array containing the first index for each + violation type: + 0: a joint whose dynamic-constraint topology changed + 1: a DoF whose finite-limit state changed + 2: a joint whose passive/actuated partition changed + 3: a joint with an unsupported combination of target modes + + An entry equal to the maximum of the joint and DoF counts indicates that no + violation of that type was found. + + Args: + model: The Newton model containing the updated joints to validate. + joints: The current Kamino joint model, before applying the updates. + built_limit_finite: The built finite limit state for each DoF. + violations: The array to store the violations. + check_dof: Whether to check the DoF updates. + check_actuation: Whether to check the actuation updates. + + Returns: + The sentinel value indicating no violations. + """ + dim = max(model.joint_count, model.joint_dof_count) + violations.fill_(dim) + if check_dof and dim > 0: + wp.launch( + kernel=validate_joint_dof_updates_kernel, + dim=dim, + inputs=[ + # Inputs: + model.joint_qd_start, + model.joint_armature, + model.joint_damping, + model.joint_target_ke, + model.joint_target_kd, + joints.num_dynamic_cts, + model.joint_limit_lower, + model.joint_limit_upper, + built_limit_finite, + model.joint_count, + model.joint_dof_count, + # Outputs: + violations, + ], + device=model.device, + ) + if check_actuation and model.joint_count > 0: + wp.launch( + kernel=validate_joint_actuation_updates_kernel, + dim=model.joint_count, + inputs=[ + # Inputs: + model.joint_qd_start, + model.joint_target_mode, + joints.act_type, + # Outputs: + violations, + ], + device=model.device, + ) + + return dim + + +def convert_model_joint_actuation(model: Model, joints: JointsModel) -> None: + """Update Kamino's per-joint actuation types from Newton target modes.""" + if model.joint_count == 0: + return + wp.launch( + kernel=update_joint_actuation_kernel, + dim=model.joint_count, + inputs=[ + # Inputs: + model.joint_qd_start, + model.joint_target_mode, + # Outputs: + joints.act_type, + ], + device=model.device, + ) + + def convert_model_joint_transforms(model: Model, joints: JointsModel) -> None: """ Converts the joint model parameterization of Newton's to Kamino's format. @@ -660,6 +928,107 @@ def convert_model_joint_transforms(model: Model, joints: JointsModel) -> None: ) +def compute_material_first_shape( + geom_material: wp.array[wp.int32], + num_materials: int, +) -> wp.array[wp.int32]: + """Compute the first shape associated with each fixed material ID. + + Args: + geom_material: Material ID for each shape. + num_materials: Number of registered materials. + + Returns: + Per-material shape indices. Materials without an associated shape use + the shape count as a sentinel. + """ + shape_count = geom_material.shape[0] + first_shape = wp.full(num_materials, shape_count, dtype=wp.int32, device=geom_material.device) + if shape_count > 0: + wp.launch( + kernel=material_first_shape_kernel, + dim=shape_count, + inputs=[ + # Inputs: + geom_material, + # Outputs: + first_shape, + ], + device=geom_material.device, + ) + return first_shape + + +def convert_model_materials( + model: Model, + model_kamino: ModelKamino, + first_shape: wp.array[wp.int32], + conflict: wp.array[wp.int32], +) -> None: + """Update Kamino's material properties in place from Newton shape materials. + + Recomputes per-material friction and restitution from + ``model.shape_material_mu`` and ``model.shape_material_restitution`` while + preserving the material arrays referenced by Kamino's collision detector. + + Args: + model: Newton model containing the updated shape materials. + model_kamino: Kamino model whose material tables are updated. + first_shape: Cached first shape associated with each fixed material ID. + conflict: Scratch scalar for reporting conflicting material updates. + + Raises: + RuntimeError: If shapes assigned to the same material ID have different + material properties and would require splitting that material. + """ + materials = model_kamino.materials + conflict.fill_(materials.num_materials) + + # Check each shape against the cached representative for its material. + wp.launch( + kernel=validate_material_update_kernel, + dim=model.shape_count, + inputs=[ + # Inputs: + model.shape_material_mu, + model.shape_material_restitution, + model_kamino.geoms.material, + first_shape, + # Outputs: + conflict, + ], + device=model.device, + ) + + conflict_material = int(conflict.numpy()[0]) + if conflict_material < materials.num_materials: + raise RuntimeError( + f"Multiple shapes assigned to contact material {conflict_material} attempted to update it with " + "different friction or restitution values; recreate SolverKamino to split the material." + ) + + # Once conflicts have been ruled out, update the material properties in place. + wp.launch( + kernel=update_materials_kernel, + dim=materials.num_materials, + inputs=[ + # Inputs: + model.shape_material_mu, + model.shape_material_restitution, + first_shape, + model.shape_count, + # Outputs: + materials.restitution, + materials.static_friction, + materials.dynamic_friction, + model_kamino.material_pairs.restitution, + model_kamino.material_pairs.static_friction, + model_kamino.material_pairs.dynamic_friction, + ], + device=model.device, + ) + + def convert_rigid_bodies( model: Model, model_size: SizeKamino, @@ -972,27 +1341,21 @@ def convert_joints( joint_dof_type_np = joint_dof_type.numpy() # Assign base bodies based on articulation roots (if articulations are present) + world_has_non_floating_root = np.zeros((model.world_count,), dtype=bool) if model.articulation_count > 0: articulation_start_np = model.articulation_start.numpy() articulation_world_np = model.articulation_world.numpy() - # For each articulation, assign its base body and joint to the corresponding world, - # if the base joint is a unary free joint. - # NOTE: We only assign the first articulation found in each world - has_non_free_root = False + # NOTE: We only assign the first articulation rooted by a unary free joint in each world for aid in range(model.articulation_count): wid = articulation_world_np[aid] base_joint = articulation_start_np[aid] base_body = joint_child_np[base_joint] if base_body_idx_np[wid] == -1 and base_joint_idx_np[wid] == -1: - if joint_dof_type_np[base_joint] != JointDoFType.FREE: - has_non_free_root = True + if joint_dof_type_np[base_joint] != JointDoFType.FREE or joint_parent_np[base_joint] != -1: + world_has_non_floating_root[wid] = True continue base_body_idx_np[wid] = base_body base_joint_idx_np[wid] = base_joint - if has_non_free_root: - msg.warning( - "Model has articulations with a non-free joint as root, disabling floating base resets for those worlds." - ) # For worlds without articulations, look for a unary free joint, or use the first body for wid in range(model.world_count): @@ -1015,6 +1378,13 @@ def convert_joints( continue base_body_idx_np[wid] = body_world_start_np[wid] + # Warn user if an articulation root couldn't be used as base because it is not a free joint + if np.any(world_has_non_floating_root & (base_body_idx_np == -1)): + msg.warning( + "Model has articulations whose root is not a free joint attached to the world, " + "disabling floating base resets for those worlds." + ) + # Update size object model_size.sum_of_num_joints = int(num_joints_np.sum()) model_size.max_of_num_joints = int(num_joints_np.max()) diff --git a/newton/_src/solvers/kamino/_src/core/gravity.py b/newton/_src/solvers/kamino/_src/core/gravity.py index 22362eeb0a..4a18e6e4cf 100644 --- a/newton/_src/solvers/kamino/_src/core/gravity.py +++ b/newton/_src/solvers/kamino/_src/core/gravity.py @@ -1,247 +1,81 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers # SPDX-License-Identifier: Apache-2.0 -"""The gravity descriptor and model used throughout Kamino""" +"""Gravity containers used by Kamino.""" from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field import numpy as np import warp as wp -from .....core.types import override +from .....core.types import Axis, override from .....sim.model import Model -from ..utils import logger as msg +from ....coupled.model_view import ModelView from .types import ArrayLike, Descriptor -### -# Module interface -### +__all__ = ["GRAVITY_DEFAULT", "GravityDescriptor", "GravityModel"] -__all__ = [ - "GRAVITY_ACCEL_DEFAULT", - "GRAVITY_DIREC_DEFAULT", - "GRAVITY_NAME_DEFAULT", - "GravityDescriptor", - "GravityModel", - "convert_model_gravity", -] +GRAVITY_DEFAULT = -9.81 +"""Default gravity along the world's up axis [m/s²].""" -### -# Module configs -### -wp.set_module_options({"enable_backward": False}) - - -### -# Constants -### - -GRAVITY_NAME_DEFAULT = "Earth" -"""The default gravity descriptor name, set as 'Earth'.""" - -GRAVITY_ACCEL_DEFAULT = 9.8067 -""" -The default gravitational acceleration in m/s^2. -Equal to Earth's standard gravity of approximately 9.8067 m/s^2. -""" +@dataclass +class GravityDescriptor(Descriptor): + """Describe a world's gravity vector.""" -GRAVITY_DIREC_DEFAULT = [0.0, 0.0, -1.0] -"""The default direction of gravity defined as -Z.""" + vector: wp.vec3f = field(default_factory=lambda: GravityDescriptor.default_from_up_axis(Axis.Z).vector) + """Gravity vector [m/s²].""" + @staticmethod + def default_from_up_axis(up_axis: Axis, *, name: str = "gravity") -> GravityDescriptor: + """Return Newton's default gravity along the negative up axis.""" + vector = wp.vec3f(*(component * GRAVITY_DEFAULT for component in up_axis.to_vector())) + return GravityDescriptor(name=name, vector=vector) -### -# Containers -### + @staticmethod + def from_array(vector: ArrayLike, *, name: str = "gravity") -> GravityDescriptor: + """Create a gravity descriptor from a three-component vector.""" + components = np.asarray(vector, dtype=np.float32) + if components.shape != (3,): + raise ValueError(f"Gravity vector must have shape (3,), got {components.shape}.") + return GravityDescriptor(name=name, vector=wp.vec3f(*components)) + @staticmethod + def from_usd( + direction: ArrayLike, magnitude: float, up_axis: Axis, distance_unit: float, *, name: str = "gravity" + ) -> GravityDescriptor: + """Create a gravity descriptor from OpenUSD scene attributes.""" + direction_array = np.asarray(direction, dtype=np.float32) + if direction_array.shape != (3,): + raise ValueError(f"Gravity direction must have shape (3,), got {direction_array.shape}.") + + direction_length = np.linalg.norm(direction_array) + if direction_length == 0.0: + direction_array = -np.asarray(up_axis.to_vector(), dtype=np.float32) + else: + direction_array /= direction_length -class GravityDescriptor(Descriptor): - """ - A container to describe a world's gravity. - - Attributes: - name: The name of the gravity descriptor. - uid: The unique identifier of the gravity descriptor. - enabled: Whether gravity is enabled. - acceleration: The gravitational acceleration magnitude [m/s²]. - direction: The normalized direction vector of gravity. - """ - - def __init__( - self, - enabled: bool = True, - acceleration: float = GRAVITY_ACCEL_DEFAULT, - direction: ArrayLike = GRAVITY_DIREC_DEFAULT, - name: str = GRAVITY_NAME_DEFAULT, - uid: str | None = None, - ): - """ - Initialize the gravity descriptor. - - Args: - enabled: Whether gravity is enabled. - Defaults to `True` to enable gravity by default. - acceleration: The gravitational acceleration magnitude in m/s^2. - Defaults to 9.8067 m/s^2 (Earth's gravity). - direction: The normalized direction vector of gravity. - Defaults to pointing down the -Z axis. - name: The name of the gravity descriptor. - uid: Optional unique identifier of the gravity descriptor. - """ - super().__init__(name, uid) - self._enabled: bool = enabled - self._acceleration: float = acceleration - self._direction: wp.vec3f = wp.normalize(wp.vec3f(direction)) + if magnitude == -float("inf"): + magnitude = abs(GRAVITY_DEFAULT) + return GravityDescriptor.from_array(direction_array * (distance_unit * magnitude), name=name) @override - def __repr__(self): - """Returns a human-readable string representation of the GravityDescriptor.""" - return ( - f"GravityDescriptor(\n" - f"name={self.name},\n" - f"uid={self.uid},\n" - f"enabled={self.enabled},\n" - f"acceleration={self.acceleration},\n" - f"direction={self.direction}\n" - f")" - ) - - @property - def enabled(self) -> bool: - """Returns whether gravity is enabled.""" - return self._enabled - - @enabled.setter - def enabled(self, on: bool): - """Sets whether gravity is enabled.""" - self._enabled = on - - @property - def acceleration(self) -> float: - """Returns the gravitational acceleration.""" - return self._acceleration - - @acceleration.setter - def acceleration(self, g: float): - """Sets the gravitational acceleration.""" - self._acceleration = g - - @property - def direction(self) -> wp.vec3f: - """Returns the normalized direction vector of gravity.""" - return self._direction - - @direction.setter - def direction(self, direction: wp.vec3f): - """Sets the normalized direction vector of gravity.""" - self._direction = wp.normalize(direction) - - def dir_accel(self) -> wp.vec4f: - """Returns the gravity direction and acceleration as compactly as a :class:`wp.vec4f`.""" - return wp.vec4f([self.direction[0], self.direction[1], self.direction[2], self.acceleration]) - - def vector(self) -> wp.vec4f: - """Returns the effective gravity vector and enabled flag compactly as a :class:`wp.vec4f`.""" - g = wp.vec3f(self.acceleration * self.direction) - return wp.vec4f([g[0], g[1], g[2], float(self.enabled)]) + def __repr__(self) -> str: + """Return a human-readable representation.""" + return f"GravityDescriptor(name={self.name!r}, uid={self.uid!r}, vector={self.vector})" @dataclass class GravityModel: - """ - A container to hold the time-invariant gravity model data. - - Attributes: - g_dir_acc: The gravity direction and acceleration vector as ``[g_dir_x, g_dir_y, g_dir_z, g_accel]``. - Shape of ``(num_worlds,)``. - vector: The gravity vector defined as ``[g_x, g_y, g_z, enabled]``. - Shape of ``(num_worlds,)``. - """ - - g_dir_acc: wp.array[wp.vec4f] | None = None - """ - The gravity direction and acceleration vector. - Shape of ``(num_worlds,)``. - """ - - vector: wp.array[wp.vec4f] | None = None - """ - The gravity vector defined as ``[g_x, g_y, g_z, enabled]``. - Shape of ``(num_worlds,)``. - """ - - ### - # Operations - ### + """Hold per-world gravity vectors.""" - @staticmethod - def from_newton(model_in: Model) -> GravityModel: - return convert_model_gravity(model_in) - - -### -# Utilities -### - - -# TODO: Re-implement using kernels -def convert_model_gravity(model_in: Model, gravity_out: GravityModel | None = None) -> GravityModel: - """ - Converts the gravity representation from the Newton model to the Kamino format. - - Args: - model_in: The input Newton model containing the gravity information to be converted. - gravity_out: The output GravityModel instance where the converted gravity data will be stored. - If `None`, a new GravityModel instance will be created and returned. - If the arrays within `gravity_out` are not already allocated - with the appropriate shapes, this function will allocate them. - """ - # Capture the necessary properties from source model - gravity_np = model_in.gravity.numpy().copy() - - # Allocate data for the conversion - g_dir_acc_np = np.zeros((model_in.world_count, 4), dtype=np.float32) - vector_np = np.zeros((model_in.world_count, 4), dtype=np.float32) - - # Convert each world's gravity vector into direction - # and acceleration, and pack into the output arrays - for w in range(model_in.world_count): - g_vec = gravity_np[w, :] - accel = float(np.linalg.norm(g_vec)) - if accel > 0.0: - direction = g_vec / accel - else: - direction = np.array([0.0, 0.0, -1.0]) - g_dir_acc_np[w, :3] = direction - g_dir_acc_np[w, 3] = accel - vector_np[w, :3] = g_vec - vector_np[w, 3] = 1.0 if accel > 0.0 else 0.0 - - # If the output gravity model is not provided, create a new one with allocated arrays; - if gravity_out is None: - with wp.ScopedDevice(model_in.device): - gravity_out = GravityModel( - g_dir_acc=wp.array(g_dir_acc_np, dtype=wp.vec4f), - vector=wp.array(vector_np, dtype=wp.vec4f), - ) - - # Otherwise, ensure the provided model has allocated arrays of the - # correct shape and type, and copy the converted data into them. - else: - # Ensure that the output GravityModel has allocated arrays of the correct shape and type - if gravity_out.g_dir_acc is None or gravity_out.g_dir_acc.shape != (model_in.world_count,): - msg.warning("Output `GravityModel.g_dir_acc` array does not have matching shape. Allocating a new array.") - gravity_out.g_dir_acc = wp.array(g_dir_acc_np, dtype=wp.vec4f, device=model_in.device) - else: - gravity_out.g_dir_acc.assign(g_dir_acc_np) - if gravity_out.vector is None or gravity_out.vector.shape != (model_in.world_count,): - msg.warning("Output `GravityModel.vector` array does not have matching shape. Allocating a new array.") - gravity_out.vector = wp.array(vector_np, dtype=wp.vec4f, device=model_in.device) - else: - gravity_out.vector.assign(vector_np) + vector: wp.array[wp.vec3] | None = None + """Per-world gravity vector [m/s²]. Shape of ``(num_worlds,)``.""" - # Return the output gravity model - return gravity_out + @staticmethod + def from_newton(model: Model | ModelView) -> GravityModel: + """Create a gravity model that aliases Newton's gravity array.""" + return GravityModel(vector=model.gravity) diff --git a/newton/_src/solvers/kamino/_src/core/joints.py b/newton/_src/solvers/kamino/_src/core/joints.py index 8cdf6db829..79d51e3185 100644 --- a/newton/_src/solvers/kamino/_src/core/joints.py +++ b/newton/_src/solvers/kamino/_src/core/joints.py @@ -5,7 +5,6 @@ from __future__ import annotations -import math from dataclasses import dataclass, field from enum import IntEnum @@ -15,7 +14,7 @@ from .....core.types import MAXVAL, override from .....sim import JointTargetMode, JointType -from .math import FLOAT32_MAX, FLOAT32_MIN, PI, TWO_PI +from .math import FLOAT32_MAX from .types import ( ArrayLike, Descriptor, @@ -225,11 +224,11 @@ def bound(self) -> float: Returns the numerical bound imposed by the correction mode. """ if self.value == self.TWOPI: - return float(TWO_PI) + return float(wp.tau) # Note: wp.tau is 2 * pi elif self.value == self.CONTINUOUS: return float(JOINT_QMAX) elif self.value == self.NONE: - return float(PI) + return float(wp.pi) else: raise ValueError(f"Unknown joint correction mode: {self.value}") @@ -304,6 +303,8 @@ class JointDoFType(IntEnum): Conventions: - Each joint connects a Base body `B` to a Follower body `F`. - The relative motion of body `F' w.r.t. body `B` defines the positive direction of the joint's DoFs. + - Mixed linear/angular vectors follow Newton's ``(linear, angular)`` ordering; translational entries + before rotational entries. - `R_x`, `R_y`, `R_z`: denote rotational DoFs about the local x, y, z axes respectively. - `T_x`, `T_y`, `T_z`: denote translational DoFs along the local x, y, z axes respectively. - Joints are indexed by `j`, and we often employ the subscript notation `*_j`. @@ -315,13 +316,13 @@ class JointDoFType(IntEnum): FREE = 0 """ - A 6-DoF free-floating joint, with rotational + translational DoFs - along {`R_x`, `R_y`, `R_z`, `T_x`, `T_y`, `T_z`}. + A 6-DoF free-floating joint, with translational + rotational DoFs + along {`T_x`, `T_y`, `T_z`, `R_x`, `R_y`, `R_z`}. Coordinates: 7D transform: 3D position + 4D unit quaternion DoFs: - 6D twist: 3D angular velocity + 3D linear velocity + 6D twist: 3D linear velocity + 3D angular velocity Constraints: None """ @@ -352,12 +353,12 @@ class JointDoFType(IntEnum): CYLINDRICAL = 3 """ - A 2-DoF cylindrical joint, with rotational + translational DoFs along {`R_x`, `T_x`}. + A 2-DoF cylindrical joint, with translational + rotational DoFs along {`T_x`, `R_x`}. Coordinates: - 2D vector of angle {`R_x`} + 1D distance {`T_x`} + 2D vector of distance {`T_x`} + angle {`R_x`} DoFs: - 2D vector of angular velocity {`R_x`} + linear velocity {`T_x`} + 2D vector of linear velocity {`T_x`} + angular velocity {`R_x`} """ # TODO: Add support for PLANAR joints with 2D linear DOFS along {`T_x`, `T_y`} @@ -443,7 +444,7 @@ def num_coords(self) -> int: elif self.value == self.PRISMATIC: return 1 # 1D distance elif self.value == self.CYLINDRICAL: - return 2 # 2D vector of angle + distance + return 2 # 2D vector of distance + angle elif self.value == self.UNIVERSAL: return 2 # 2D angles elif self.value == self.SPHERICAL: @@ -461,13 +462,13 @@ def num_dofs(self) -> int: Returns the number of DoFs defined by the joint DoF type. """ if self.value == self.FREE: - return 6 # 3D angular velocity + 3D linear velocity + return 6 # 3D linear velocity + 3D angular velocity elif self.value == self.REVOLUTE: return 1 # 1D angular velocity elif self.value == self.PRISMATIC: return 1 # 1D linear velocity elif self.value == self.CYLINDRICAL: - return 2 # 1D angular velocity + 1D linear velocity + return 2 # 1D linear velocity + 1D angular velocity elif self.value == self.UNIVERSAL: return 2 # 2D angular velocities elif self.value == self.SPHERICAL: @@ -883,7 +884,7 @@ def num_coords_wp(dof_type: int) -> int: elif dof_type == JointDoFType.PRISMATIC: return 1 # 1D distance elif dof_type == JointDoFType.CYLINDRICAL: - return 2 # 2D vector of angle + distance + return 2 # 2D vector of distance + angle elif dof_type == JointDoFType.UNIVERSAL: return 2 # 2D angles elif dof_type == JointDoFType.SPHERICAL: @@ -908,13 +909,13 @@ def num_dofs_wp(dof_type: int) -> int: invalid. """ if dof_type == JointDoFType.FREE: - return 6 # 3D angular velocity + 3D linear velocity + return 6 # 3D linear velocity + 3D angular velocity elif dof_type == JointDoFType.REVOLUTE: return 1 # 1D angular velocity elif dof_type == JointDoFType.PRISMATIC: return 1 # 1D linear velocity elif dof_type == JointDoFType.CYLINDRICAL: - return 2 # 1D angular velocity + 1D linear velocity + return 2 # 1D linear velocity + 1D angular velocity elif dof_type == JointDoFType.UNIVERSAL: return 2 # 2D angular velocities elif dof_type == JointDoFType.SPHERICAL: @@ -1586,12 +1587,7 @@ def _check_dofs_array( return [float(default) for _ in range(size)] if isinstance(x, (int, float, np.floating)): - if x == math.inf: - return [float(FLOAT32_MAX) for _ in range(size)] - elif x == -math.inf: - return [float(FLOAT32_MIN) for _ in range(size)] - else: - return [x] * size + return [x] * size if isinstance(x, ArrayLike): if len(x) == 0: @@ -1601,11 +1597,6 @@ def _check_dofs_array( raise ValueError(f"Invalid DOF array length: {len(x)} != {size}") if all(isinstance(x, (float, np.floating)) for x in x): - for i in range(len(x)): - if x[i] == math.inf: - x[i] = float(FLOAT32_MAX) - elif x[i] == -math.inf: - x[i] = float(FLOAT32_MIN) return x else: raise TypeError(f"Unsupported DOF array type: {type(x)!r}; expected float, iterable of floats, or None") @@ -1753,9 +1744,10 @@ class JointsModel: """ Minimum (a.k.a. lower) joint DoF limits of each joint (as flat array). - Limits are dimensioned according to the number of DoFs of each joint, - as opposed to the number of coordinates in order to handle cases such - where joints have more coordinates than DoFs (e.g. spherical joints). + Although applying to joint coordinates, limits are dimensioned + according to the number of DoFs of each joint, as the number of limits + depends on the intrinsic number of DoFs, not on its (possibly redundant, + e.g. for spherical joints) parameterization into coordinates. Shape of ``(sum_of_num_joint_dofs,)``. """ @@ -1764,9 +1756,10 @@ class JointsModel: """ Maximum (a.k.a. upper) joint DoF limits of each joint (as flat array). - Limits are dimensioned according to the number of DoFs of each joint, - as opposed to the number of coordinates in order to handle cases such - where joints have more coordinates than DoFs (e.g. spherical joints). + Although applying to joint coordinates, limits are dimensioned + according to the number of DoFs of each joint, as the number of limits + depends on the intrinsic number of DoFs, not on its (possibly redundant, + e.g. for spherical joints) parameterization into coordinates. Shape of ``(sum_of_num_joint_dofs,)``. """ diff --git a/newton/_src/solvers/kamino/_src/core/materials.py b/newton/_src/solvers/kamino/_src/core/materials.py index aacde7aee0..63a0dba03f 100644 --- a/newton/_src/solvers/kamino/_src/core/materials.py +++ b/newton/_src/solvers/kamino/_src/core/materials.py @@ -39,7 +39,6 @@ ### __all__ = [ - "DEFAULT_DENSITY", "DEFAULT_FRICTION", "DEFAULT_RESTITUTION", "MaterialDescriptor", @@ -54,12 +53,6 @@ # Constants ### -DEFAULT_DENSITY = 1000.0 -""" -The global default density for materials, in kg/m^3. -Equals ``1000.0`` kg/m^3. -""" - DEFAULT_RESTITUTION = 0.0 """ The global default restitution coefficient for material pairs. @@ -122,17 +115,13 @@ class MaterialDescriptor(Descriptor): """ A container to represent a managed material. - This descriptor holds both intrinsic and extrinsic properties of a material. While the former - are truly dependent on the material itself (e.g., density), the latter are actually dependent - on the pairwise interactions of the material with others (e.g., friction, restitution). These - extrinsic properties are stored here to support model specifications such as USD which - currently do not support material-pair definitions. + To support model specifications such as USD (which currently do not support material-pair + definitions), this descriptor holds extrinsic properties of a material (e.g., friction, restitution), + which are actually dependent on the pairwise interactions of the material with others. Attributes: name: The name of the material. uid: The unique identifier (UUID) of the material. - density: The density of the material [kg/m³]. - Defaults to the global default of ``1000.0`` kg/m³. restitution: The coefficient of restitution, according to the Newtonian impact model. Defaults to the global default of ``0.0``. static_friction: The coefficient of static friction, according to the Coulomb friction model. @@ -149,12 +138,6 @@ class MaterialDescriptor(Descriptor): # Attributes ### - density: float = DEFAULT_DENSITY - """ - The density of the material, in kg/m^3. - Defaults to the global default of ``1000.0`` kg/m^3. - """ - restitution: float = DEFAULT_RESTITUTION """ The coefficient of restitution, according to the Newtonian impact model. @@ -196,7 +179,6 @@ def __repr__(self) -> str: f"MaterialDescriptor(\n" f"name: {self.name},\n" f"uid: {self.uid},\n" - f"density: {self.density},\n" f"restitution: {self.restitution},\n" f"static_friction: {self.static_friction},\n" f"dynamic_friction: {self.dynamic_friction}\n" @@ -254,8 +236,6 @@ class MaterialsModel: Attributes: num_materials: Total number of materials represented in the model. - density: Array of material density values of each registered material. - Shape of ``(num_materials,)``. restitution: Array of restitution coefficients for each registered material. Shape of ``(num_materials,)``. static_friction: Array of static friction coefficients for each registered material. @@ -267,12 +247,6 @@ class MaterialsModel: num_materials: int = 0 """Total number of materials represented in the model.""" - density: wp.array[wp.float32] | None = None - """ - Array of material density values of each registered material. - Shape of ``(num_materials,)``. - """ - restitution: wp.array[wp.float32] | None = None """ Array of restitution coefficients for each registered material. diff --git a/newton/_src/solvers/kamino/_src/core/math.py b/newton/_src/solvers/kamino/_src/core/math.py index fa2db00af2..50500c4e53 100644 --- a/newton/_src/solvers/kamino/_src/core/math.py +++ b/newton/_src/solvers/kamino/_src/core/math.py @@ -9,7 +9,7 @@ import numpy as np import warp as wp -from warp._src.types import Any, Float +from warp._src.types import Any from .....core.types import Axis, AxisType from .types import ( @@ -37,53 +37,9 @@ FLOAT32_EPS = wp.constant(wp.float32(np.finfo(np.float32).eps)) """Machine epsilon for 32-bit float: the smallest value such that 1.0 + eps != 1.0.""" -UNIT_X = wp.constant(wp.vec3f(1.0, 0.0, 0.0)) -""" 3D unit vector for the X axis """ - -UNIT_Y = wp.constant(wp.vec3f(0.0, 1.0, 0.0)) -""" 3D unit vector for the Y axis """ - -UNIT_Z = wp.constant(wp.vec3f(0.0, 0.0, 1.0)) -""" 3D unit vector for the Z axis """ - -PI = wp.constant(3.141592653589793) -"""Convenience constant for PI""" - -TWO_PI = wp.constant(6.283185307179586) -"""Convenience constant for 2 * PI""" - -HALF_PI = wp.constant(1.5707963267948966) -"""Convenience constant for PI / 2""" - -COS_PI_6 = wp.constant(0.8660254037844387) -"""Convenience constant for cos(PI / 6)""" - -I_2 = wp.constant(wp.mat22f(1, 0, 0, 1)) -""" The 2x2 identity matrix.""" - I_3 = wp.constant(wp.mat33f(1, 0, 0, 0, 1, 0, 0, 0, 1)) """ The 3x3 identity matrix.""" -I_4 = wp.constant(wp.mat44f(1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1)) -""" The 4x4 identity matrix.""" - -I_6 = wp.constant( - wp.spatial_matrixf( - 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1 - ) -) -""" The 6x6 identity matrix.""" - - -### -# General-purpose functions -### - - -@wp.func -def squared_norm(x: Any) -> Float: - return wp.dot(x, x) - ### # Rotation matrices @@ -108,68 +64,6 @@ def axis_to_mat33(axis: AxisType) -> wp.mat33f: return wp.mat33f(0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0) -@wp.func -def R_x(theta: wp.float32) -> wp.mat33f: - """ - Computes the rotation matrix around the X axis. - - Args: - theta: The angle in radians. - - Returns: - The rotation matrix. - """ - c = wp.cos(theta) - s = wp.sin(theta) - return wp.mat33f(1.0, 0.0, 0.0, 0.0, c, -s, 0.0, s, c) - - -@wp.func -def R_y(theta: wp.float32) -> wp.mat33f: - """ - Computes the rotation matrix around the Y axis. - - Args: - theta: The angle in radians. - - Returns: - The rotation matrix. - """ - c = wp.cos(theta) - s = wp.sin(theta) - return wp.mat33f(c, 0.0, s, 0.0, 1.0, 0.0, -s, 0.0, c) - - -@wp.func -def R_z(theta: wp.float32) -> wp.mat33f: - """ - Computes the rotation matrix around the Z axis. - - Args: - theta: The angle in radians. - - Returns: - The rotation matrix. - """ - c = wp.cos(theta) - s = wp.sin(theta) - return wp.mat33f(c, -s, 0.0, s, c, 0.0, 0.0, 0.0, 1.0) - - -@wp.func -def unskew(S: wp.mat33f) -> wp.vec3f: - """ - Extracts the 3D vector from a 3x3 skew-symmetric matrix. - - Args: - S: The 3x3 skew-symmetric matrix. - - Returns: - The vector extracted from the skew-symmetric matrix. - """ - return wp.vec3f(S[2, 1], S[0, 2], S[1, 0]) - - ### # Quaternions ### @@ -229,31 +123,6 @@ def H_of(q: wp.quatf) -> mat34f: return H -@wp.func -def quat_from_vec4(v: wp.vec4f) -> wp.quatf: - """ - Convert a wp.vec4f to a quaternion type. - """ - return wp.quatf(v[0], v[1], v[2], v[3]) - - -@wp.func -def quat_to_vec4(q: wp.quatf) -> wp.vec4f: - """ - Convert a quaternion type to a wp.vec4f. - """ - return wp.vec4f(q.x, q.y, q.z, q.w) - - -@wp.func -def quat_conj(q: wp.quatf) -> wp.quatf: - """ - Compute the conjugate of a quaternion. - The conjugate of a quaternion q = (x, y, z, w) is defined as: q_conj = (x, y, z, -w) - """ - return wp.quatf(q.x, q.y, q.z, -q.w) - - @wp.func def quat_positive(q: wp.quatf) -> wp.quatf: """ @@ -276,36 +145,6 @@ def quat_imaginary(q: wp.quatf) -> wp.vec3f: return wp.vec3f(q.x, q.y, q.z) -@wp.func -def quat_apply(q: wp.quatf, v: wp.vec3f) -> wp.vec3f: - """ - Apply a quaternion to a vector. - The quaternion is applied to the vector using the formula: - v' = s * v + q.w * uv + qv x uv, where s = ||q||^2, uv = 2 * qv x v, and qv is the imaginary part of the quaternion. - """ - qv = quat_imaginary(q) - uv = 2.0 * wp.cross(qv, v) - s = wp.dot(q, q) - return s * v + q.w * uv + wp.cross(qv, uv) - - -@wp.func -def quat_derivative(q: wp.quatf, omega: wp.vec3f) -> wp.quatf: - """ - Computes the quaternion derivative from a quaternion and angular velocity. - - Args: - q: The quaternion of the current pose of the body. - omega: The angular velocity of the body. - - Returns: - The quaternion derivative. - """ - vdq = 0.5 * wp.transpose(G_of(q)) * omega - dq = wp.quaternion(vdq.x, vdq.y, vdq.z, vdq.w, dtype=wp.float32) - return dq - - @wp.func def quat_log(q: wp.quatf) -> wp.vec3f: """ @@ -334,39 +173,6 @@ def quat_log(q: wp.quatf) -> wp.vec3f: return c * pv -@wp.func -def quat_log_decomposed(q: wp.quatf) -> wp.vec4f: - """ - Computes the logarithm of a quaternion using the stable - `4 * atan()` formulation to render an angle-axis vector. - - The output is a wp.vec4f with the following format: - - `a = [x, y, z, c]` is the angle-axis output - - `[x, y, z]` is the axis of rotation - - `c` is the angle. - """ - p = quat_positive(q) - pv = quat_imaginary(p) - pv_norm_sq = wp.dot(pv, pv) - pw_sq = p.w * p.w - pv_norm = wp.sqrt(pv_norm_sq) - - # Check if the norm of the imaginary part is infinitesimal - if pv_norm_sq > FLOAT32_EPS: - # Regular solution for larger angles - # Use more stable 4 * atan() formulation over the 2 * atan(pv_norm / pw) - # TODO: angle = 4.0 * wp.atan2(pv_norm, (p.w + wp.sqrt(pw_sq + pv_norm_sq))) - angle = 4.0 * wp.atan(pv_norm / (p.w + wp.sqrt(pw_sq + pv_norm_sq))) - c = angle / pv_norm - else: - # Taylor expansion solution for small angles - # For the alternative branch use the limit of angle / pv_norm for angle -> 0.0 - c = (2.0 - wp.static(2.0 / 3.0) * (pv_norm_sq / pw_sq)) / p.w - - # Return the scaled imaginary part of the quaternion - return wp.vec4f(pv.x, pv.y, pv.z, c) - - @wp.func def quat_exp(v: wp.vec3f) -> wp.quatf: """ @@ -400,92 +206,6 @@ def quat_exp(v: wp.vec3f) -> wp.quatf: return q -@wp.func -def quat_product(q1: wp.quatf, q2: wp.quatf) -> wp.quatf: - """ - Computes the quaternion product of two quaternions. - - Args: - q1: The first quaternion. - q2: The second quaternion. - - Returns: - The result of the quaternion product. - """ - q3 = wp.quat_identity(dtype=wp.float32) - q3.x = q1.w * q2.x + q1.x * q2.w + q1.y * q2.z - q1.z * q2.y - q3.y = q1.w * q2.y - q1.x * q2.z + q1.y * q2.w + q1.z * q2.x - q3.z = q1.w * q2.z + q1.x * q2.y - q1.y * q2.x + q1.z * q2.w - q3.w = q1.w * q2.w - q1.x * q2.x - q1.y * q2.y - q1.z * q2.z - return q3 - - -@wp.func -def quat_box_plus(q: wp.quatf, v: wp.vec3f) -> wp.quatf: - """ - Computes the box-plus operation for a quaternion and a vector: - R(q) [+] v == exp(v) * R(q), where R(q) is the rotation matrix of the quaternion q. - - Args: - q: The quaternion. - v: The vector. - - Returns: - The result of the box-plus operation. - """ - return quat_product(quat_exp(v), q) - - -@wp.func -def quat_from_x_rot(angle_rad: wp.float32) -> wp.quatf: - """ - Computes a unit quaternion corresponding to rotation by given angle about the x axis - """ - return wp.quatf(wp.sin(0.5 * angle_rad), 0.0, 0.0, wp.cos(0.5 * angle_rad)) - - -@wp.func -def quat_from_y_rot(angle_rad: wp.float32) -> wp.quatf: - """ - Computes a unit quaternion corresponding to rotation by given angle about the y axis - """ - return wp.quatf(0.0, wp.sin(0.5 * angle_rad), 0.0, wp.cos(0.5 * angle_rad)) - - -@wp.func -def quat_from_z_rot(angle_rad: wp.float32) -> wp.quatf: - """ - Computes a unit quaternion corresponding to rotation by given angle about the z axis - """ - return wp.quatf(0.0, 0.0, wp.sin(0.5 * angle_rad), wp.cos(0.5 * angle_rad)) - - -@wp.func -def quat_to_euler_xyz(q: wp.quatf) -> wp.vec3f: - """ - Converts a unit quaternion to XYZ Euler angles (also known as Cardan angles). - """ - rpy = wp.vec3f(0.0) - R_20 = -2.0 * (q.x * q.z - q.w * q.y) - if wp.abs(R_20) < 1.0: - rpy[1] = wp.asin(-R_20) - rpy[0] = wp.atan2(2.0 * (q.y * q.z + q.w * q.x), q.w * q.w - q.x * q.x - q.y * q.y + q.z * q.z) - rpy[2] = wp.atan2(2.0 * (q.x * q.y + q.w * q.z), q.w * q.w + q.x * q.x - q.y * q.y - q.z * q.z) - else: # Gimbal lock - rpy[0] = wp.atan2(-2.0 * (q.x * q.y - q.w * q.z), q.w * q.w - q.x * q.x + q.y * q.y - q.z * q.z) - rpy[1] = wp.half_pi if R_20 <= -1.0 else -wp.half_pi - rpy[2] = 0.0 - return rpy - - -@wp.func -def quat_from_euler_xyz(rpy: wp.vec3f) -> wp.quatf: - """ - Converts XYZ Euler angles (also known as Cardan angles) to a unit quaternion. - """ - return wp.quat_from_matrix(R_z(rpy.z) @ R_y(rpy.y) @ R_x(rpy.x)) - - @wp.func def quat_left_jacobian_inverse(q: wp.quatf) -> wp.mat33f: """ @@ -510,28 +230,6 @@ def quat_left_jacobian_inverse(q: wp.quatf) -> wp.mat33f: return wp.identity(3, dtype=wp.float32) - wp.skew(c0 * pv) + wp.skew(c1 * pv) * wp.skew(pv) -@wp.func -def quat_normalized_apply(q: wp.quatf, v: wp.vec3f) -> wp.vec3f: - """ - Combines quaternion normalization and applying a unit quaternion to a vector - """ - qv = quat_imaginary(q) - s = wp.dot(q, q) - uv_s = (2.0 / s) * wp.cross(qv, v) - return v + q[3] * uv_s + wp.cross(qv, uv_s) - - -@wp.func -def quat_conj_normalized_apply(q: wp.quatf, v: wp.vec3f) -> wp.vec3f: - """ - Combines quaternion conjugation, normalization and applying a unit quaternion to a vector - """ - qv = quat_imaginary(q) - s = wp.dot(q, q) - uv_s = (2.0 / s) * wp.cross(qv, v) - return v - q[3] * uv_s + wp.cross(qv, uv_s) - - @wp.func def quat_twist_angle(q: wp.quatf, axis: wp.vec3f) -> wp.float32: """ @@ -680,49 +378,6 @@ def unit_quat_conj_apply_jacobian(q: wp.quatf, v: wp.vec3f) -> mat34f: ### -@wp.func -def screw(linear: wp.vec3f, angular: wp.vec3f) -> wp.spatial_vectorf: - """ - Constructs a 6D screw (as `wp.spatial_vectorf`) from 3D linear and angular components. - - Args: - linear: The linear component of the screw. - angular: The angular component of the screw. - - Returns: - The resulting screw represented as a 6D vector. - """ - return wp.spatial_vectorf(linear[0], linear[1], linear[2], angular[0], angular[1], angular[2]) - - -@wp.func -def screw_linear(s: wp.spatial_vectorf) -> wp.vec3f: - """ - Extracts the linear component from a 6D screw vector. - - Args: - s: The 6D screw vector. - - Returns: - The linear component of the screw. - """ - return wp.vec3f(s[0], s[1], s[2]) - - -@wp.func -def screw_angular(s: wp.spatial_vectorf) -> wp.vec3f: - """ - Extracts the angular component from a 6D screw vector. - - Args: - s: The 6D screw vector. - - Returns: - The angular component of the screw. - """ - return wp.vec3f(s[3], s[4], s[5]) - - @wp.func def screw_transform_matrix_from_points(r_A: wp.vec3f, r_B: wp.vec3f) -> wp.spatial_matrixf: """ @@ -747,7 +402,7 @@ def screw_transform_matrix_from_points(r_A: wp.vec3f, r_B: wp.vec3f) -> wp.spati The 6x6 screw transformation matrix. """ # Initialize the wrench matrix - W_BA = I_6 + W_BA = wp.identity(n=6, dtype=wp.float32) # Fill the lower left block with the skew-symmetric matrix S_BA = wp.skew(r_A - r_B) @@ -870,11 +525,11 @@ def compute_body_twist_update_with_eom( w_i: wp.spatial_vectorf, ) -> tuple[wp.vec3f, wp.vec3f]: # Extract linear and angular parts - v_i = screw_linear(u_i) - omega_i = screw_angular(u_i) + v_i = wp.spatial_top(u_i) + omega_i = wp.spatial_bottom(u_i) S_i = wp.skew(omega_i) - f_i = screw_linear(w_i) - tau_i = screw_angular(w_i) + f_i = wp.spatial_top(w_i) + tau_i = wp.spatial_bottom(w_i) # Compute velocity update equations v_i_n = v_i + dt * (g + inv_m_i * f_i) @@ -897,7 +552,7 @@ def compute_body_pose_update_with_logmap( # Compute configuration update equations r_i_n = r_i + dt * v_i - q_i_n = quat_box_plus(q_i, dt * omega_i) + q_i_n = quat_exp(dt * omega_i) * q_i p_i_n = wp.transformf(r_i_n, q_i_n) # Return the new pose and twist diff --git a/newton/_src/solvers/kamino/_src/core/shapes.py b/newton/_src/solvers/kamino/_src/core/shapes.py index 101aca192a..eba835b177 100644 --- a/newton/_src/solvers/kamino/_src/core/shapes.py +++ b/newton/_src/solvers/kamino/_src/core/shapes.py @@ -620,7 +620,7 @@ def max_contacts_for_shape_pair(type_a: int, type_b: int) -> tuple[int, int]: Count the number of potential contact points for a collision pair in both directions of the collision pair (collisions from A to B and from B to A). - Inputs must be canonicalized such that the type of shape A is less than or equal to the type of shape B. + Shape types are canonicalized such that the type of shape A is less than or equal to the type of shape B. Args: type_a: First shape type as :class:`GeoType` integer value. @@ -633,6 +633,19 @@ def max_contacts_for_shape_pair(type_a: int, type_b: int) -> tuple[int, int]: if type_a > type_b: type_a, type_b = type_b, type_a + return _max_contacts_for_shape_pair_impl(type_a, type_b) + + +@wp.func +def _max_contacts_for_shape_pair_impl(type_a: int, type_b: int) -> tuple[int, int]: + """ + Return the contact capacity for a canonical shape pair without reordering. + + Noncanonical pairs return ``(0, 0)``. + """ + if type_a > type_b: + return 0, 0 + if type_a == GeoType.SPHERE: return 1, 0 @@ -649,8 +662,6 @@ def max_contacts_for_shape_pair(type_a: int, type_b: int) -> tuple[int, int]: return _MESH_CONVEX_MAX, 0 elif type_b == GeoType.CONE: return 4, 4 - elif type_b == GeoType.PLANE: - return 8, 8 elif type_a == GeoType.ELLIPSOID: if type_b == GeoType.ELLIPSOID: @@ -663,8 +674,6 @@ def max_contacts_for_shape_pair(type_a: int, type_b: int) -> tuple[int, int]: return _MESH_CONVEX_MAX, 0 elif type_b == GeoType.CONE: return 8, 8 - elif type_b == GeoType.PLANE: - return 4, 4 elif type_a == GeoType.CYLINDER: if type_b == GeoType.CYLINDER: @@ -675,8 +684,6 @@ def max_contacts_for_shape_pair(type_a: int, type_b: int) -> tuple[int, int]: return _MESH_CONVEX_MAX, 0 elif type_b == GeoType.CONE: return 4, 4 - elif type_b == GeoType.PLANE: - return 6, 6 elif type_a == GeoType.BOX: if type_b == GeoType.BOX: @@ -685,16 +692,12 @@ def max_contacts_for_shape_pair(type_a: int, type_b: int) -> tuple[int, int]: return _MESH_CONVEX_MAX, 0 elif type_b == GeoType.CONE: return 8, 8 - elif type_b == GeoType.PLANE: - return 12, 12 elif type_a == GeoType.MESH or type_a == GeoType.CONVEX_MESH: if type_b == GeoType.HFIELD: return _MESH_MESH_MAX, 0 elif type_b == GeoType.CONE: return _MESH_CONVEX_MAX, 0 - elif type_b == GeoType.PLANE: - return _MESH_CONVEX_MAX, 0 else: return _MESH_MESH_MAX, 0 @@ -705,11 +708,24 @@ def max_contacts_for_shape_pair(type_a: int, type_b: int) -> tuple[int, int]: elif type_a == GeoType.CONE: if type_b == GeoType.CONE: return 4, 4 - elif type_b == GeoType.PLANE: - return 8, 8 elif type_a == GeoType.PLANE: - pass + if type_b == GeoType.HFIELD: + return _MESH_CONVEX_MAX, 0 + elif type_b == GeoType.SPHERE: + return 1, 0 + elif type_b == GeoType.CAPSULE: + return 8, 8 + elif type_b == GeoType.ELLIPSOID: + return 4, 4 + elif type_b == GeoType.CYLINDER: + return 6, 6 + elif type_b == GeoType.BOX: + return 12, 12 + elif type_b == GeoType.MESH or type_b == GeoType.CONVEX_MESH: + return _MESH_CONVEX_MAX, 0 + elif type_b == GeoType.CONE: + return 8, 8 # unsupported type combination return 0, 0 diff --git a/newton/_src/solvers/kamino/_src/core/types.py b/newton/_src/solvers/kamino/_src/core/types.py index 4f9b6df933..a008f99167 100644 --- a/newton/_src/solvers/kamino/_src/core/types.py +++ b/newton/_src/solvers/kamino/_src/core/types.py @@ -6,7 +6,7 @@ from __future__ import annotations import uuid -from collections.abc import Iterable +from collections.abc import Sequence from dataclasses import dataclass import numpy as np @@ -28,7 +28,7 @@ IntType = wp.int16 | wp.int32 | wp.int64 VecIntType = wp.vec2s | wp.vec2i | wp.vec2l -ArrayLike = np.ndarray | list | tuple | Iterable +ArrayLike = np.ndarray | list | tuple | Sequence """An Array-like structure for aliasing various data types compatible with numpy.""" diff --git a/newton/_src/solvers/kamino/_src/dynamics/delassus.py b/newton/_src/solvers/kamino/_src/dynamics/delassus.py index 31fb0a2f69..a84ed619d1 100644 --- a/newton/_src/solvers/kamino/_src/dynamics/delassus.py +++ b/newton/_src/solvers/kamino/_src/dynamics/delassus.py @@ -1755,10 +1755,33 @@ def device(self) -> wp.DeviceLike: def constraint_jacobian(self) -> BlockSparseMatrices[wp.float32, wp.int32, vec6f]: return self._jacobians._J_cts.bsm + @property + def regularization(self) -> wp.array[wp.float32]: + """Active diagonal regularization used by sparse matrix-vector products.""" + regularization = self._combined_regularization + if regularization is None: + regularization = self._eta + if regularization is None: + raise RuntimeError("Sparse Delassus regularization has not been configured.") + return regularization + ### # Operations ### + def apply_jacobian_transpose( + self, + x: wp.array[wp.float32], + y: wp.array[wp.float32], + world_mask: wp.array[wp.bool], + ) -> None: + """Apply the current transposed constraint Jacobian to a vector.""" + if self.ATy_op is None or self._transpose_op_matrix is None: + raise RuntimeError("Sparse Delassus transpose operator has not been assigned.") + if self._needs_update: + self.update() + self.ATy_op(self._transpose_op_matrix, x, y, world_mask) + def matvec(self, x: wp.array[wp.float32], y: wp.array[wp.float32], world_mask: wp.array[wp.bool]): """ Performs the sparse matrix-vector product `y = D @ x`, applying regularization and diff --git a/newton/_src/solvers/kamino/_src/dynamics/dual.py b/newton/_src/solvers/kamino/_src/dynamics/dual.py index 07c4b85ccd..52b7c51257 100644 --- a/newton/_src/solvers/kamino/_src/dynamics/dual.py +++ b/newton/_src/solvers/kamino/_src/dynamics/dual.py @@ -61,7 +61,7 @@ from .....core.types import override from ...config import ConfigBase, ConstrainedDynamicsConfig, ConstraintStabilizationConfig from ..core.data import DataKamino -from ..core.math import FLOAT32_EPS, screw, screw_angular, screw_linear +from ..core.math import FLOAT32_EPS from ..core.model import ModelKamino from ..core.size import SizeKamino from ..core.types import vec6f @@ -343,7 +343,7 @@ def gravity_plus_coriolis_wrench( """ f_gi_i = m_i * g tau_gi_i = -wp.skew(omega_i) @ (I_i @ omega_i) - return screw(f_gi_i, tau_gi_i) + return wp.spatial_vectorf(*f_gi_i, *tau_gi_i) @wp.func @@ -370,7 +370,7 @@ def gravity_plus_coriolis_wrench_split( def _build_nonlinear_generalized_force( # Inputs: model_time_dt: wp.array[wp.float32], - model_gravity_vector: wp.array[wp.vec4f], + model_gravity_vector: wp.array[wp.vec3f], model_bodies_wid: wp.array[wp.int32], model_bodies_m_i: wp.array[wp.float32], state_bodies_u_i: wp.array[wp.spatial_vectorf], @@ -393,13 +393,10 @@ def _build_nonlinear_generalized_force( # Get world data dt = model_time_dt[wid] - gv = model_gravity_vector[wid] - - # Extract the effective gravity vector - g = gv.w * wp.vec3f(gv.x, gv.y, gv.z) + g = model_gravity_vector[wid] # Extract the linear and angular components of the generalized velocity - omega_i = screw_angular(u_i) + omega_i = wp.spatial_bottom(u_i) # Compute the net external wrench on the body h_i = w_e_i + w_a_i + gravity_plus_coriolis_wrench(g, m_i, I_i, omega_i) @@ -412,7 +409,7 @@ def _build_nonlinear_generalized_force( def _build_generalized_free_velocity( # Inputs: model_time_dt: wp.array[wp.float32], - model_gravity_vector: wp.array[wp.vec4f], + model_gravity_vector: wp.array[wp.vec3f], model_bodies_wid: wp.array[wp.int32], model_bodies_m_i: wp.array[wp.float32], model_bodies_inv_m_i: wp.array[wp.float32], @@ -439,26 +436,23 @@ def _build_generalized_free_velocity( # Get world data dt = model_time_dt[wid] - gv = model_gravity_vector[wid] - - # Extract the effective gravity vector - g = gv.w * wp.vec3f(gv.x, gv.y, gv.z) + g = model_gravity_vector[wid] # Extract the linear and angular components of the generalized velocity - v_i = screw_linear(u_i) - omega_i = screw_angular(u_i) + v_i = wp.spatial_top(u_i) + omega_i = wp.spatial_bottom(u_i) # Compute the net external wrench on the body h_i = w_e_i + w_a_i + gravity_plus_coriolis_wrench(g, m_i, I_i, omega_i) - f_h_i = screw_linear(h_i) - tau_h_i = screw_angular(h_i) + f_h_i = wp.spatial_top(h_i) + tau_h_i = wp.spatial_bottom(h_i) # Compute the generalized free-velocity vector components v_f_i = v_i + dt * (inv_m_i * f_h_i) omega_f_i = omega_i + dt * (inv_I_i @ tau_h_i) # Store the generalized free-velocity vector - problem_u_f[bid] = screw(v_f_i, omega_f_i) + problem_u_f[bid] = wp.spatial_vectorf(*v_f_i, *omega_f_i) @wp.kernel diff --git a/newton/_src/solvers/kamino/_src/geometry/contacts.py b/newton/_src/solvers/kamino/_src/geometry/contacts.py index 995c3b10f9..6166e3e5f4 100644 --- a/newton/_src/solvers/kamino/_src/geometry/contacts.py +++ b/newton/_src/solvers/kamino/_src/geometry/contacts.py @@ -37,7 +37,6 @@ from .....sim.model import Model from .....sim.state import State from ..core.materials import MaterialMixMode, make_get_mixed_material_pair_property -from ..core.math import COS_PI_6, UNIT_X, UNIT_Y from ..core.model import ModelKamino from ..core.types import ( to_warp_int32_array, @@ -378,14 +377,16 @@ def reset(self): # Functions ### +COS_PI_6 = wp.constant(0.8660254037844387) + @wp.func def make_contact_frame_znorm(n: wp.vec3f) -> wp.mat33f: n = wp.normalize(n) - if wp.abs(wp.dot(n, UNIT_X)) < COS_PI_6: - e = UNIT_X + if wp.abs(n[0]) < COS_PI_6: + e = wp.vec3f(1.0, 0.0, 0.0) else: - e = UNIT_Y + e = wp.vec3f(0.0, 1.0, 0.0) o = wp.normalize(wp.cross(n, e)) t = wp.normalize(wp.cross(o, n)) return wp.mat33f(t.x, o.x, n.x, t.y, o.y, n.y, t.z, o.z, n.z) @@ -394,10 +395,10 @@ def make_contact_frame_znorm(n: wp.vec3f) -> wp.mat33f: @wp.func def make_contact_frame_xnorm(n: wp.vec3f) -> wp.mat33f: n = wp.normalize(n) - if wp.abs(wp.dot(n, UNIT_X)) < COS_PI_6: - e = UNIT_X + if wp.abs(n[0]) < COS_PI_6: + e = wp.vec3f(1.0, 0.0, 0.0) else: - e = UNIT_Y + e = wp.vec3f(0.0, 1.0, 0.0) o = wp.normalize(wp.cross(n, e)) t = wp.normalize(wp.cross(o, n)) return wp.mat33f(n.x, t.x, o.x, n.y, t.y, o.y, n.z, t.z, o.z) @@ -622,7 +623,7 @@ def gapfunc(self) -> wp.array[wp.vec4f]: The ``w`` component stores the signed ``distance`` between margin-shifted surfaces: - ``w < 0`` means penetration past the resting separation defined by the margin - - ``w > 0`` means separation within the detection ``distance = gap + margin`` + - ``w > 0`` means margin-shifted separation within the detection gap """ self._assert_has_data() return self._data.gapfunc @@ -1351,7 +1352,7 @@ def convert_contacts_newton_to_kamino( # convert exceeds the capacity of the output contacts. if contacts_in.rigid_contact_max > contacts_out.model_max_contacts_host: msg.warning( - "Newton `rigid_contact_max` (%d) exceeds Kamino `model_max_contacts_host` (%d); contacts will be truncated.", + "Newton `rigid_contact_max` (%d) exceeds Kamino `model_max_contacts_host` (%d); active contacts may be truncated.", contacts_in.rigid_contact_max, contacts_out.model_max_contacts_host, ) @@ -1359,9 +1360,10 @@ def convert_contacts_newton_to_kamino( # Skip conversion of contact forces if not requested contacts_in_force = contacts_in.force if convert_forces else None - # Set the maximum number of contacts to convert to the smallest of the - # number of contacts detected and the maximum capacity of the output contacts. - max_converted_contacts = min(contacts_in.rigid_contact_max, contacts_out.model_max_contacts_host) + # Scan every Newton contact slot so saturated worlds cannot prevent later + # worlds from filling their own capacity. The kernel skips inactive slots + # and enforces per-world and model limits. + max_converted_contacts = contacts_in.rigid_contact_max # Clear the output contacts to reset the active contact # counts and reset contact data to sentinel values. @@ -1373,9 +1375,7 @@ def convert_contacts_newton_to_kamino( restitution_mix_mode=MaterialMixMode.from_string(restitution_mix_mode), ) - # Launch the conversion kernel to convert Newton contacts to Kamino's format - # NOTE: To reduce overhead, the total thread count is set to the smallest of - # the number of contacts detected and the maximum capacity of the output contacts. + # Launch the conversion kernel to convert Newton contacts to Kamino's format. wp.launch( kernel=_convert_contacts_newton_to_kamino, dim=max_converted_contacts, @@ -1493,11 +1493,10 @@ def convert_contacts_kamino_to_newton( f"contacts_out.device={contacts_out.device}" ) - # Issue warning to the user if the number of contacts to - # convert exceeds the capacity of the output contacts. - if contacts_in.model_max_contacts_host > contacts_out.rigid_contact_max: + # Only Kamino-generated contacts can be truncated during export. + if clear_output and contacts_in.model_max_contacts_host > contacts_out.rigid_contact_max: msg.warning( - "Kamino `model_max_contacts_host` (%d) exceeds Newton `rigid_contact_max` (%d); contacts will be truncated.", + "Kamino `model_max_contacts_host` (%d) exceeds Newton `rigid_contact_max` (%d); active contacts may be truncated.", contacts_in.model_max_contacts_host, contacts_out.rigid_contact_max, ) diff --git a/newton/_src/solvers/kamino/_src/geometry/detector.py b/newton/_src/solvers/kamino/_src/geometry/detector.py index 074133f4af..62b7573bd8 100644 --- a/newton/_src/solvers/kamino/_src/geometry/detector.py +++ b/newton/_src/solvers/kamino/_src/geometry/detector.py @@ -147,6 +147,86 @@ def __repr__(self): return self.__str__() +### +# Contact capacity helpers +### + +# Conservative heuristics for the fallback allocation path when pair-based +# capacity metadata is unavailable (``model_minimum_contacts == 0``). +_EXPLICIT_CONTACTS_PER_PAIR = 10 +_DYNAMIC_CONTACTS_PER_COLLIDABLE = 20 + + +def _cap_world_contacts_at_total(world_max_contacts: list[int], max_total: int) -> list[int]: + """Scale per-world contact budgets down so their sum does not exceed ``max_total``.""" + total = sum(world_max_contacts) + if total <= max_total: + return list(world_max_contacts) + if max_total <= 0: + return [0] * len(world_max_contacts) + + capped = [0] * len(world_max_contacts) + remainders: list[tuple[float, int]] = [] + assigned = 0 + for i, count in enumerate(world_max_contacts): + scaled = count * max_total / total + floor = int(scaled) + capped[i] = floor + assigned += floor + remainders.append((scaled - floor, i)) + for _, i in sorted(remainders, key=lambda item: item[0], reverse=True): + if assigned >= max_total: + break + capped[i] += 1 + assigned += 1 + return capped + + +def _estimate_fallback_world_max_contacts( + model: ModelKamino, + config: CollisionDetectorConfig, +) -> list[int]: + """Estimate per-world contact capacity from geometry when pair metadata is unavailable.""" + num_worlds = model.size.num_worlds + world_max_contacts = [0] * num_worlds + + if config.broadphase == "explicit" and model.geoms.collidable_pairs is not None: + pairs = model.geoms.collidable_pairs.numpy() + wid = model.geoms.wid.numpy() + for pair in pairs: + g0, g1 = int(pair[0]), int(pair[1]) + world_id = int(wid[g0]) if wid[g0] >= 0 else int(wid[g1]) + if 0 <= world_id < num_worlds: + world_max_contacts[world_id] += _EXPLICIT_CONTACTS_PER_PAIR + else: + wid = model.geoms.wid.numpy() + group = model.geoms.group.numpy() + for geom_id in range(len(wid)): + world_id = int(wid[geom_id]) + if 0 <= world_id < num_worlds and group[geom_id] > 0: + world_max_contacts[world_id] += _DYNAMIC_CONTACTS_PER_COLLIDABLE + + return world_max_contacts + + +def _resolve_contact_capacity( + model: ModelKamino, + config: CollisionDetectorConfig, +) -> tuple[int, list[int]]: + """Resolve model- and per-world contact budgets from geometry and config caps.""" + if model.geoms.model_minimum_contacts > 0: + world_max_contacts = list(model.geoms.world_minimum_contacts) + else: + world_max_contacts = _estimate_fallback_world_max_contacts(model, config) + + model_max_contacts = sum(world_max_contacts) + if model_max_contacts > config.max_contacts: + world_max_contacts = _cap_world_contacts_at_total(world_max_contacts, config.max_contacts) + model_max_contacts = sum(world_max_contacts) + + return model_max_contacts, world_max_contacts + + ### # Interfaces ### @@ -303,41 +383,22 @@ def finalize( # Configure the collision detection pipeline type based on the config self._pipeline_type = CollisionPipelineType.from_string(self._config.pipeline) - # TODO: FIX THIS SO THAT PER-WORLD MAX IS ACTUALLY BASED ON THE NUM OF COLLIDABLE - # GOEMS IN EACH WORLD, INSTEAD OF JUST DIVIDING THE MODEL MAX BY THE NUM WORLDS - # For collision pipeline, we don't multiply by per-pair factors since broad phase - # discovers pairs dynamically. Users can provide rigid_contact_max explicitly, - # otherwise it is estimated from shape count and broad phase mode. - if self._model.geoms.model_minimum_contacts > 0: - self._model_max_contacts = self._model.geoms.model_minimum_contacts - self._world_max_contacts = self._model.geoms.world_minimum_contacts - else: - # Estimate based on broad phase mode and available information - if self._config.broadphase == "explicit" and self._model.geoms.collidable_pairs is not None: - # For EXPLICIT mode, we know the maximum possible pairs - # Estimate ~10 contacts per shape pair (conservative for mesh-mesh contacts) - self._model_max_contacts = max(self._config.max_contacts, self._model.geoms.num_collidable_pairs * 10) - else: - # For NXN/SAP dynamic broad phase, estimate based on shape count - # Assume each shape contacts ~20 others on average (conservative estimate) - # This scales much better than O(N²) while still being safe - self._model_max_contacts = max(self._config.max_contacts, self._model.geoms.num_collidable * 20) - - # Set the world max contacts to be the same for all worlds in the model - num_worlds = self._model.size.num_worlds - self._world_max_contacts = [self._model_max_contacts // num_worlds] * num_worlds - - # Override per-world max contacts if config specifies it. + # Resolve contact capacity. if self._config.max_contacts_per_world is not None: + # Use the explicit per-world override when available. num_worlds = self._model.size.num_worlds per_world = self._config.max_contacts_per_world self._world_max_contacts = [per_world] * num_worlds self._model_max_contacts = per_world * num_worlds + else: + # Otherwise estimate per world from geometry. + # ``max_contacts`` caps the model total. + self._model_max_contacts, self._world_max_contacts = _resolve_contact_capacity(self._model, self._config) # Create the contacts interface which will allocate all contacts data arrays # NOTE: If internal allocations happen, then they will contain # the contacts generated by the collision detection pipelines - self._contacts = ContactsKamino(capacity=self._world_max_contacts, device=self._device) + self._contacts = ContactsKamino(capacity=list(self._world_max_contacts), device=self._device) # Proceed with allocations only if the model admits contacts, which # occurs when collision geometries defined in the builder and model diff --git a/newton/_src/solvers/kamino/_src/integrators/euler.py b/newton/_src/solvers/kamino/_src/integrators/euler.py index 001b4f0c36..e52e0b9418 100644 --- a/newton/_src/solvers/kamino/_src/integrators/euler.py +++ b/newton/_src/solvers/kamino/_src/integrators/euler.py @@ -14,11 +14,7 @@ from .....core.types import override from ..core.control import ControlKamino from ..core.data import DataKamino -from ..core.math import ( - compute_body_pose_update_with_logmap, - compute_body_twist_update_with_eom, - screw, -) +from ..core.math import compute_body_pose_update_with_logmap, compute_body_twist_update_with_eom from ..core.model import ModelKamino from ..core.state import StateKamino from ..geometry.contacts import ContactsKamino @@ -82,7 +78,7 @@ def euler_semi_implicit_with_logmap( ) # Return the new pose and twist - return p_i_n, screw(v_i_n, omega_i_n) + return p_i_n, wp.spatial_vectorf(*v_i_n, *omega_i_n) ### @@ -95,7 +91,7 @@ def _integrate_semi_implicit_euler_inplace( # Inputs: alpha: float, model_dt: wp.array[wp.float32], - model_gravity: wp.array[wp.vec4f], + model_gravity: wp.array[wp.vec3f], model_bodies_wid: wp.array[wp.int32], model_bodies_inv_m: wp.array[wp.float32], model_bodies_I: wp.array[wp.mat33f], @@ -113,8 +109,7 @@ def _integrate_semi_implicit_euler_inplace( # Retrieve the time step and gravity vector dt = model_dt[wid] - gv = model_gravity[wid] - g = gv.w * wp.vec3f(gv.x, gv.y, gv.z) + g = model_gravity[wid] # Retrieve the model data inv_m_i = model_bodies_inv_m[tid] diff --git a/newton/_src/solvers/kamino/_src/integrators/moreau.py b/newton/_src/solvers/kamino/_src/integrators/moreau.py index 08520b9aa8..cf27987009 100644 --- a/newton/_src/solvers/kamino/_src/integrators/moreau.py +++ b/newton/_src/solvers/kamino/_src/integrators/moreau.py @@ -15,13 +15,7 @@ from .....core.types import override from ..core.control import ControlKamino from ..core.data import DataKamino -from ..core.math import ( - compute_body_pose_update_with_logmap, - compute_body_twist_update_with_eom, - screw, - screw_angular, - screw_linear, -) +from ..core.math import compute_body_pose_update_with_logmap, compute_body_twist_update_with_eom from ..core.model import ModelKamino from ..core.state import StateKamino from ..geometry.contacts import ContactsKamino @@ -85,7 +79,7 @@ def moreau_jean_semi_implicit_with_logmap( ) # Return the new pose and twist - return p_i_n, screw(v_i_n, omega_i_n) + return p_i_n, wp.spatial_vectorf(*v_i_n, *omega_i_n) ### @@ -119,8 +113,8 @@ def _integrate_moreau_jean_first_inplace( q_i_m = compute_body_pose_update_with_logmap( dt=0.5 * dt, p_i=q_i, - v_i=screw_linear(u_i), - omega_i=screw_angular(u_i), + v_i=wp.spatial_top(u_i), + omega_i=wp.spatial_bottom(u_i), ) # Store the computed next pose and twist @@ -132,7 +126,7 @@ def _integrate_moreau_jean_second_inplace( # Inputs: alpha: float, model_dt: wp.array[wp.float32], - model_gravity: wp.array[wp.vec4f], + model_gravity: wp.array[wp.vec3f], model_bodies_wid: wp.array[wp.int32], model_bodies_inv_m: wp.array[wp.float32], model_bodies_I: wp.array[wp.mat33f], @@ -151,8 +145,7 @@ def _integrate_moreau_jean_second_inplace( # Retrieve the configured time-step and the # gravity vector of the corresponding world dt = model_dt[wid] - gv = model_gravity[wid] - g = gv.w * wp.vec3f(gv.x, gv.y, gv.z) + g = model_gravity[wid] # Retrieve the model data inv_m_i = model_bodies_inv_m[tid] diff --git a/newton/_src/solvers/kamino/_src/kinematics/jacobians.py b/newton/_src/solvers/kamino/_src/kinematics/jacobians.py index 22ff9be904..99ec4cf128 100644 --- a/newton/_src/solvers/kamino/_src/kinematics/jacobians.py +++ b/newton/_src/solvers/kamino/_src/kinematics/jacobians.py @@ -59,6 +59,65 @@ ### +@wp.func +def build_full_joint_jacobian( + joint_id: wp.int32, + dof_type: wp.int32, + bid_B: wp.int32, + bid_F: wp.int32, + model_joints_X_Bj: wp.array[wp.mat33f], + model_joints_X_Fj: wp.array[wp.mat33f], + state_joints_p: wp.array[wp.transformf], + state_bodies_q: wp.array[wp.transformf], +): + """ + Computes the full (6x6) joint Jacobian (constraint and DoFs) for a specific joint. + """ + # Retrieve the pose transform of the joint + T_j = state_joints_p[joint_id] + r_j = wp.transform_get_translation(T_j) + R_X_j = wp.quat_to_matrix(wp.transform_get_rotation(T_j)) + + # Retrieve the pose transforms of each body + T_B_j = wp.transform_identity() + if bid_B > -1: + T_B_j = state_bodies_q[bid_B] + T_F_j = state_bodies_q[bid_F] + r_B_j = wp.transform_get_translation(T_B_j) + r_F_j = wp.transform_get_translation(T_F_j) + + if dof_type == JointDoFType.FREE: + # By Newton's convention, a free joint's twist and wrench (specified in `joint_qd` and + # `joint_f`) are both defined at the child's center of mass, with world-aligned axes. Since + # Kamino avoids a conversion, the change of reference frame needs to be taken into account + # in the Jacobian. + JT_F_j = wp.identity(n=6, dtype=wp.float32) + JT_B_j = -screw_transform_matrix_from_points(r_F_j, r_B_j) + + else: + # Compute the wrench matrices + # TODO: Since the lever-arm is a relative position, can we just use B_r_Bj and F_r_Fj instead? + W_j_B = screw_transform_matrix_from_points(r_j, r_B_j) + W_j_F = screw_transform_matrix_from_points(r_j, r_F_j) + + # General case: Compute the effective projector to joint frame and expand to 6D + if dof_type != JointDoFType.UNIVERSAL: + R_X_bar_j = expand6d(R_X_j) + # Universal joint: replace R_X_j with the frame of the intermediate body for rotation constraints + else: + j_q_j = compute_joint_relative_quaternion( + T_B_j, T_F_j, model_joints_X_Bj[joint_id], model_joints_X_Fj[joint_id] + ) + R_intermediate = compute_intermediate_body_frame_universal_joint(j_q_j) + R_X_bar_j = concat6d(R_X_j, R_X_j @ R_intermediate) + + # Compute the extended jacobians, i.e. without the selection-matrix multiplication + JT_B_j = -W_j_B @ R_X_bar_j # Reaction is on the Base body body ; (6 x 6) + JT_F_j = W_j_F @ R_X_bar_j # Action is on the Follower body ; (6 x 6) + + return JT_B_j, JT_F_j + + def make_store_joint_jacobian_dense_func(axes: Any): """ Generates a warp function to store body-pair Jacobian blocks into a target flat @@ -472,36 +531,17 @@ def _build_joint_jacobians_dense( J_jdc_row_start = J_cjmio + nbd * (jdcgo + dyn_cts_offset_world) J_jkc_row_start = J_cjmio + nbd * (jkcgo + kin_cts_offset_world) - # Retrieve the pose transform of the joint - T_j = state_joints_p[jid] - r_j = wp.transform_get_translation(T_j) - R_X_j = wp.quat_to_matrix(wp.transform_get_rotation(T_j)) - - # Retrieve the pose transforms of each body - T_B_j = wp.transform_identity() - if bid_B > -1: - T_B_j = state_bodies_q[bid_B] - T_F_j = state_bodies_q[bid_F] - r_B_j = wp.transform_get_translation(T_B_j) - r_F_j = wp.transform_get_translation(T_F_j) - - # Compute the wrench matrices - # TODO: Since the lever-arm is a relative position, can we just use B_r_Bj and F_r_Fj instead? - W_j_B = screw_transform_matrix_from_points(r_j, r_B_j) - W_j_F = screw_transform_matrix_from_points(r_j, r_F_j) - - # General case: Compute the effective projector to joint frame and expand to 6D - if dof_type != JointDoFType.UNIVERSAL: - R_X_bar_j = expand6d(R_X_j) - # Universal joint: replace R_X_j with the frame of the intermediate body for rotation constraints - else: - j_q_j = compute_joint_relative_quaternion(T_B_j, T_F_j, model_joints_X_Bj[jid], model_joints_X_Fj[jid]) - R_intermediate = compute_intermediate_body_frame_universal_joint(j_q_j) - R_X_bar_j = concat6d(R_X_j, R_X_j @ R_intermediate) - - # Compute the extended jacobians, i.e. without the selection-matrix multiplication - JT_B_j = -W_j_B @ R_X_bar_j # Reaction is on the Base body body ; (6 x 6) - JT_F_j = W_j_F @ R_X_bar_j # Action is on the Follower body ; (6 x 6) + # Compute the full jacobians, i.e. without the selection-matrix multiplication + JT_B_j, JT_F_j = build_full_joint_jacobian( + jid, + dof_type, + bid_B, + bid_F, + model_joints_X_Bj, + model_joints_X_Fj, + state_joints_p, + state_bodies_q, + ) # Store joint dynamic constraint jacobians if applicable # NOTE: We use the extraction method for DoFs since dynamic constraints are in DoF-space @@ -561,36 +601,17 @@ def _build_joint_jacobians_sparse( bid_B = model_joints_bid_B[jid] bid_F = model_joints_bid_F[jid] - # Retrieve the pose transform of the joint - T_j = state_joints_p[jid] - r_j = wp.transform_get_translation(T_j) - R_X_j = wp.quat_to_matrix(wp.transform_get_rotation(T_j)) - - # Retrieve the pose transforms of each body - T_B_j = wp.transform_identity() - if bid_B > -1: - T_B_j = state_bodies_q[bid_B] - T_F_j = state_bodies_q[bid_F] - r_B_j = wp.transform_get_translation(T_B_j) - r_F_j = wp.transform_get_translation(T_F_j) - - # Compute the wrench matrices - # TODO: Since the lever-arm is a relative position, can we just use B_r_Bj and F_r_Fj instead? - W_j_B = screw_transform_matrix_from_points(r_j, r_B_j) - W_j_F = screw_transform_matrix_from_points(r_j, r_F_j) - - # General case: Compute the effective projector to joint frame and expand to 6D - if dof_type != JointDoFType.UNIVERSAL: - R_X_bar_j = expand6d(R_X_j) - # Universal joint: replace R_X_j with the frame of the intermediate body for rotation constraints - else: - j_q_j = compute_joint_relative_quaternion(T_B_j, T_F_j, model_joints_X_Bj[jid], model_joints_X_Fj[jid]) - R_intermediate = compute_intermediate_body_frame_universal_joint(j_q_j) - R_X_bar_j = concat6d(R_X_j, R_X_j @ R_intermediate) - - # Compute the extended jacobians, i.e. without the selection-matrix multiplication - JT_B_j = -W_j_B @ R_X_bar_j # Reaction is on the Base body body ; (6 x 6) - JT_F_j = W_j_F @ R_X_bar_j # Action is on the Follower body ; (6 x 6) + # Compute the full jacobians, i.e. without the selection-matrix multiplication + JT_B_j, JT_F_j = build_full_joint_jacobian( + jid, + dof_type, + bid_B, + bid_F, + model_joints_X_Bj, + model_joints_X_Fj, + state_joints_p, + state_bodies_q, + ) # Store joint dynamic constraint jacobians if applicable # NOTE: We use the extraction method for DoFs since dynamic constraints are in DoF-space @@ -1701,6 +1722,21 @@ def finalize( self._J_dofs_joint_nzb_offsets = to_warp_int32_array(J_dofs_joint_nzb_offsets, device=device) self._J_cts_num_joint_nzb = to_warp_int32_array(J_cts_nnzb_min, device=device) + @property + def joint_constraint_nzb_count(self) -> wp.array[wp.int32]: + """Number of joint-constraint blocks in each world.""" + return self._J_cts_num_joint_nzb + + @property + def limit_constraint_nzb_offsets(self) -> wp.array[wp.int32]: + """Global sparse-block offsets for each limit constraint.""" + return self._J_cts_limit_nzb_offsets + + @property + def contact_constraint_nzb_offsets(self) -> wp.array[wp.int32]: + """Global sparse-block offsets for each contact constraint.""" + return self._J_cts_contact_nzb_offsets + def build( self, model: ModelKamino, diff --git a/newton/_src/solvers/kamino/_src/kinematics/joints.py b/newton/_src/solvers/kamino/_src/kinematics/joints.py index 59d9153d05..f6febe0046 100644 --- a/newton/_src/solvers/kamino/_src/kinematics/joints.py +++ b/newton/_src/solvers/kamino/_src/kinematics/joints.py @@ -13,17 +13,7 @@ from ..core.data import DataKamino from ..core.joints import JointActuationType, JointCorrectionMode, JointDoFType -from ..core.math import ( - FLOAT32_MAX, - TWO_PI, - quat_log, - quat_to_vec4, - quat_twist_angle, - screw, - screw_angular, - screw_linear, - squared_norm, -) +from ..core.math import FLOAT32_MAX, quat_log, quat_twist_angle from ..core.model import ModelKamino from ..core.types import ( vec1f, @@ -66,15 +56,19 @@ @wp.func -def correct_rotational_coord( - q_j_in: wp.float32, q_j_ref: wp.float32 = 0.0, q_j_limit: wp.float32 = FLOAT32_MAX -) -> wp.float32: +def correct_rotational_coord(q_j_in: wp.float32, q_j_ref: wp.float32 = 0.0) -> wp.float32: """ Corrects a rotational joint coordinate to be as close as possible to a reference coordinate. """ - q_j_in += wp.round((q_j_ref - q_j_in) / TWO_PI) * TWO_PI - q_j_in = wp.mod(q_j_in, q_j_limit) - return q_j_in + return q_j_in + wp.round((q_j_ref - q_j_in) / wp.tau) * wp.tau # Note: wp.tau is 2 * pi + + +@wp.func +def correct_rotational_coord_with_limit( + q_j_in: wp.float32, q_j_ref: wp.float32 = 0.0, q_j_limit: wp.float32 = FLOAT32_MAX +) -> wp.float32: + """Corrects a rotational coordinate relative to a reference and wraps it within a limit.""" + return wp.mod(correct_rotational_coord(q_j_in, q_j_ref), q_j_limit) @wp.func @@ -86,7 +80,7 @@ def correct_quat_vector_coord(q_j_in: wp.vec4f, q_j_ref: wp.vec4f) -> wp.vec4f: closer to the reference quaternion `q_j_ref`, accounting for the fact that quaternions `q` and `-q` represent the same rotation. """ - if squared_norm(q_j_in + q_j_ref) < squared_norm(q_j_in - q_j_ref): + if wp.length_sq(q_j_in + q_j_ref) < wp.length_sq(q_j_in - q_j_ref): q_j_in *= -1.0 return q_j_in @@ -101,7 +95,7 @@ def correct_joint_coord_free(q_j_in: vec7f, q_j_ref: vec7f, q_j_limit: vec7f = D @wp.func def correct_joint_coord_revolute(q_j_in: vec1f, q_j_ref: vec1f, q_j_limit: vec1f = DEFAULT_LIMIT_V1F) -> vec1f: """Corrects the rotational joint coordinate.""" - q_j_in[0] = correct_rotational_coord(q_j_in[0], q_j_ref[0], q_j_limit[0]) + q_j_in[0] = correct_rotational_coord_with_limit(q_j_in[0], q_j_ref[0], q_j_limit[0]) return q_j_in @@ -116,7 +110,7 @@ def correct_joint_coord_cylindrical( q_j_in: wp.vec2f, q_j_ref: wp.vec2f, q_j_limit: wp.vec2f = DEFAULT_LIMIT_V2F ) -> wp.vec2f: """Corrects only the rotational joint coordinate.""" - q_j_in[1] = correct_rotational_coord(q_j_in[1], q_j_ref[1], q_j_limit[1]) + q_j_in[1] = correct_rotational_coord_with_limit(q_j_in[1], q_j_ref[1], q_j_limit[1]) return q_j_in @@ -125,8 +119,8 @@ def correct_joint_coord_universal( q_j_in: wp.vec2f, q_j_ref: wp.vec2f, q_j_limit: wp.vec2f = DEFAULT_LIMIT_V2F ) -> wp.vec2f: """Corrects each of the two rotational joint coordinates individually.""" - q_j_in[0] = correct_rotational_coord(q_j_in[0], q_j_ref[0], q_j_limit[0]) - q_j_in[1] = correct_rotational_coord(q_j_in[1], q_j_ref[1], q_j_limit[1]) + q_j_in[0] = correct_rotational_coord_with_limit(q_j_in[0], q_j_ref[0], q_j_limit[0]) + q_j_in[1] = correct_rotational_coord_with_limit(q_j_in[1], q_j_ref[1], q_j_limit[1]) return q_j_in @@ -227,7 +221,7 @@ def map_to_joint_coords_universal(j_r_j: wp.vec3f, j_q_j: wp.quatf) -> wp.vec2f: @wp.func def map_to_joint_coords_spherical(j_r_j: wp.vec3f, j_q_j: wp.quatf) -> wp.vec4f: """Returns the 4D unit-quaternion representing the joint rotation.""" - return quat_to_vec4(j_q_j) + return wp.vec4f(*j_q_j) @wp.func @@ -344,8 +338,10 @@ def convert_angular_vel_to_universal_joint_intermediary_frame( a_z = wp.cross(a_x, a_y) # Project angular velocity into intermediary body frame - omega = screw_angular(j_u_j) - return screw(screw_linear(j_u_j), wp.vec3f(wp.dot(omega, a_x), wp.dot(omega, a_y), wp.dot(omega, a_z))) + omega = wp.spatial_bottom(j_u_j) + return wp.spatial_vectorf( + *wp.spatial_top(j_u_j), *wp.vec3f(wp.dot(omega, a_x), wp.dot(omega, a_y), wp.dot(omega, a_z)) + ) ### @@ -400,7 +396,7 @@ def _write_typed_joint_data( if wp.static(num_cts > 0): # Construct a 6D residual vector j_theta_j = wp.static(get_joint_constraint_angular_residual_function(dof_type))(j_q_j) - j_p_j = screw(j_r_j, j_theta_j) + j_p_j = wp.spatial_vectorf(*j_r_j, *j_theta_j) # Store the joint constraint residuals for j in range(num_cts): r_j_out[cts_offset + j] = j_p_j[cts_axes[j]] @@ -636,14 +632,14 @@ def compute_joint_pose_and_relative_motion( # Extract the decomposed state of the Base body r_B_j = wp.transform_get_translation(T_B_j) q_B_j = wp.transform_get_rotation(T_B_j) - v_B_j = screw_linear(u_B_j) - omega_B_j = screw_angular(u_B_j) + v_B_j = wp.spatial_top(u_B_j) + omega_B_j = wp.spatial_bottom(u_B_j) # Extract the decomposed state of the Follower body r_F_j = wp.transform_get_translation(T_F_j) q_F_j = wp.transform_get_rotation(T_F_j) - v_F_j = screw_linear(u_F_j) - omega_F_j = screw_angular(u_F_j) + v_F_j = wp.spatial_top(u_F_j) + omega_F_j = wp.spatial_bottom(u_F_j) # Local joint frame quantities r_Bj = wp.quat_rotate(q_B_j, B_r_Bj) @@ -666,7 +662,7 @@ def compute_joint_pose_and_relative_motion( # TODO: How can we simplify this expression and make it more efficient? j_v_j = wp.quat_rotate_inv(q_Bj, v_F_j - v_B_j + wp.cross(omega_F_j, r_Fj) - wp.cross(omega_B_j, r_Bj + r_j)) j_omega_j = wp.quat_rotate_inv(q_Bj, omega_F_j - omega_B_j) - j_u_j = screw(j_v_j, j_omega_j) + j_u_j = wp.spatial_vectorf(*j_v_j, *j_omega_j) # Return the computed joint frame pose and relative motion vectors return p_j, j_r_j, j_q_j, j_u_j @@ -895,7 +891,7 @@ def _compute_joints_data( @wp.kernel def _extract_actuators_state_from_joints( # Inputs: - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported model_joint_wid: wp.array[wp.int32], model_joint_act_type: wp.array[wp.int32], model_joint_coords_offset: wp.array[wp.int32], @@ -916,7 +912,7 @@ def _extract_actuators_state_from_joints( act_type = model_joint_act_type[jid] # Early exit the operation if the joint's world is flagged as skipped or if the joint is not actuated - if not world_mask[wid] or act_type == JointActuationType.PASSIVE: + if (world_mask and not world_mask[wid]) or act_type == JointActuationType.PASSIVE: return # Retrieve the joint model data @@ -1063,11 +1059,11 @@ def compute_joints_data( def extract_actuators_state_from_joints( model: ModelKamino, - world_mask: wp.array[wp.bool], joint_q: wp.array[wp.float32], joint_u: wp.array[wp.float32], actuator_q: wp.array[wp.float32], actuator_u: wp.array[wp.float32], + world_mask: wp.array[wp.bool] | None = None, ): """ Extracts the states of the actuated joints from the full joint state arrays. @@ -1085,7 +1081,7 @@ def extract_actuators_state_from_joints( Shape of ``(sum_of_num_actuated_joint_coords,)``. actuator_u: The output array to store the actuated joint velocities. Shape of ``(sum_of_actuated_joint_dofs,)``. - world_mask: An array indicating which worlds are active (True) or skipped (False). + world_mask: Per-world boolean mask. If provided, indicates in which worlds to perform the operation. Shape of ``(num_worlds,)``. """ wp.launch( diff --git a/newton/_src/solvers/kamino/_src/kinematics/limits.py b/newton/_src/solvers/kamino/_src/kinematics/limits.py index 1f82dca935..fba1ce4a4e 100644 --- a/newton/_src/solvers/kamino/_src/kinematics/limits.py +++ b/newton/_src/solvers/kamino/_src/kinematics/limits.py @@ -10,11 +10,7 @@ import warp as wp from ..core.joints import JOINT_QMAX, JOINT_QMIN, JointDoFType -from ..core.math import ( - quat_from_vec4, - quat_log, - screw, -) +from ..core.math import quat_log from ..core.model import ModelKamino from ..core.types import ( to_warp_int32_array, @@ -198,8 +194,8 @@ def reset(self): @wp.func def map_joint_coords_to_dofs_free(q_j: vec7f) -> wp.spatial_vectorf: """Maps free joint quaternion to a local axes-aligned rotation vector.""" - v_j = quat_log(quat_from_vec4(q_j[3:7])) - return screw(q_j[0:3], v_j) + v_j = quat_log(wp.quatf(*q_j[3:7])) + return wp.spatial_vectorf(*q_j[0:3], *v_j) @wp.func @@ -230,7 +226,7 @@ def map_joint_coords_to_dofs_universal(q_j: wp.vec2f) -> wp.vec2f: def map_joint_coords_to_dofs_spherical(q_j: wp.vec4f) -> wp.vec3f: """Maps quaternion coordinates of a spherical joint to a local axes-aligned rotation vector.""" - return quat_log(quat_from_vec4(q_j)) + return quat_log(wp.quatf(*q_j)) @wp.func diff --git a/newton/_src/solvers/kamino/_src/kinematics/resets.py b/newton/_src/solvers/kamino/_src/kinematics/resets.py index fd3fec494d..94df52a2df 100644 --- a/newton/_src/solvers/kamino/_src/kinematics/resets.py +++ b/newton/_src/solvers/kamino/_src/kinematics/resets.py @@ -15,9 +15,10 @@ from ..kinematics.joints import ( compute_joint_pose_and_relative_motion, convert_angular_vel_to_universal_joint_intermediary_frame, + correct_quat_vector_coord, + correct_rotational_coord, get_joint_coords_mapping_function, ) -from ..solvers.fk.kernels import _correct_joint_angle, _correct_joint_quaternion ### # Module interface @@ -59,25 +60,25 @@ def _correct_joint_coords( pass # No correction needed elif wp.static(dof_type == JointDoFType.CYLINDRICAL): # Correct angle up to +/- 2 pi - coords[1] = _correct_joint_angle(coords[1], coords_ref[1]) + coords[1] = correct_rotational_coord(coords[1], coords_ref[1]) elif wp.static(dof_type == JointDoFType.FREE): # Correct quaternion up to sign quat = wp.vec4f(coords[3], coords[4], coords[5], coords[6]) quat_ref = wp.vec4f(coords_ref[3], coords_ref[4], coords_ref[5], coords_ref[6]) - quat_corrected = _correct_joint_quaternion(quat, quat_ref) + quat_corrected = correct_quat_vector_coord(quat, quat_ref) for i in range(4): coords[3 + i] = quat_corrected[i] elif wp.static(dof_type == JointDoFType.REVOLUTE): # Correct angle up to +/- 2 pi - coords[0] = _correct_joint_angle(coords[0], coords_ref[0]) + coords[0] = correct_rotational_coord(coords[0], coords_ref[0]) elif wp.static(dof_type == JointDoFType.SPHERICAL): # Correct quaternion up to sign quat_ref = wp.vec4f(coords_ref[0], coords_ref[1], coords_ref[2], coords_ref[3]) - coords = _correct_joint_quaternion(coords, quat_ref) + coords = correct_quat_vector_coord(coords, quat_ref) elif wp.static(dof_type == JointDoFType.UNIVERSAL): # Correct angles up to +/- 2 pi - coords[0] = _correct_joint_angle(coords[0], coords_ref[0]) - coords[1] = _correct_joint_angle(coords[1], coords_ref[1]) + coords[0] = correct_rotational_coord(coords[0], coords_ref[0]) + coords[1] = correct_rotational_coord(coords[1], coords_ref[1]) return coords @@ -205,7 +206,7 @@ def _get_base_q_from_joint_q_and_body_q( model_joint_coords_offset: wp.array[wp.int32], state_joint_q: wp.array[wp.float32], state_body_q: wp.array[wp.transformf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported # Outputs: base_q: wp.array[wp.transformf], ): @@ -213,7 +214,7 @@ def _get_base_q_from_joint_q_and_body_q( wid = wp.tid() # Early return based on mask - if not world_mask[wid]: + if world_mask and not world_mask[wid]: return # Read base_q from joint_q if a base joint was set for this world @@ -246,7 +247,7 @@ def _get_base_u_from_joint_u_and_body_u( model_joint_dofs_offset: wp.array[wp.int32], state_joint_u: wp.array[wp.float32], state_body_u: wp.array[wp.spatial_vectorf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported # Outputs: base_u: wp.array[wp.spatial_vectorf], ): @@ -254,7 +255,7 @@ def _get_base_u_from_joint_u_and_body_u( wid = wp.tid() # Early return based on mask - if not world_mask[wid]: + if world_mask and not world_mask[wid]: return # Read base_u from joint_u if a base joint was set for this world @@ -283,13 +284,13 @@ def _set_body_q( # Inputs: body_world_id: wp.array[wp.int32], body_q_in: wp.array[wp.transformf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported # Outputs: body_q_out: wp.array[wp.transformf], ): body_id = wp.tid() wid = body_world_id[body_id] - if not world_mask[wid]: + if world_mask and not world_mask[wid]: return body_q_out[body_id] = body_q_in[body_id] @@ -311,7 +312,7 @@ def _reset_joints_state_from_bodies_state( joint_q_0: wp.array[wp.float32], body_q: wp.array[wp.transformf], body_u: wp.array[wp.spatial_vectorf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported # Outputs joint_q: wp.array[wp.float32], joint_q_prev: wp.array[wp.float32], @@ -323,7 +324,7 @@ def _reset_joints_state_from_bodies_state( # Early return based on mask wid = joint_world_id[jid] - if not world_mask[wid]: + if world_mask and not world_mask[wid]: return # Retrieve the joint model data @@ -368,7 +369,7 @@ def _reset_joints_state_from_bodies_state( def _reset_body_velocities( # Inputs body_world_id: wp.array[wp.int32], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported # Outputs body_u: wp.array[wp.spatial_vectorf], ): @@ -377,7 +378,7 @@ def _reset_body_velocities( # Early return based on mask wid = body_world_id[body_id] - if not world_mask[wid]: + if world_mask and not world_mask[wid]: return # Reset velocities to zero @@ -388,7 +389,7 @@ def _reset_body_velocities( def _reset_body_wrenches( # Inputs body_world_id: wp.array[wp.int32], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported # Outputs body_w: wp.array[wp.spatial_vectorf], body_w_e: wp.array[wp.spatial_vectorf], @@ -398,7 +399,7 @@ def _reset_body_wrenches( # Early return based on mask wid = body_world_id[body_id] - if not world_mask[wid]: + if world_mask and not world_mask[wid]: return # Reset wrenches to zero @@ -409,7 +410,7 @@ def _reset_body_wrenches( @wp.kernel def _reset_time_of_select_worlds( # Inputs: - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported # Outputs: data_time: wp.array[wp.float32], data_steps: wp.array[wp.int32], @@ -418,7 +419,7 @@ def _reset_time_of_select_worlds( wid = wp.tid() # Skip resetting time if the world has not been marked for reset - if not world_mask[wid]: + if world_mask and not world_mask[wid]: return # Reset both the physical time and step count to zero @@ -439,7 +440,7 @@ def _eval_floating_base_relative_transform( base_u: wp.array[wp.spatial_vectorf], # None also supported body_q: wp.array[wp.transformf], body_u: wp.array[wp.spatial_vectorf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported relative_base_u: wp.bool, # Outputs: rel_transform: wp.array[wp.transformf], @@ -450,7 +451,7 @@ def _eval_floating_base_relative_transform( wid = wp.tid() # Early return based on mask - if not world_mask[wid]: + if world_mask and not world_mask[wid]: return # Determine new pose of the base body (= follower of the base joint if there is a base joint) @@ -528,7 +529,7 @@ def _apply_floating_base_transform( rel_transform: wp.array[wp.transformf], rel_velocity: wp.array[wp.spatial_vectorf], new_base_pos: wp.array[wp.vec3f], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool], # None also supported # Outputs: body_q: wp.array[wp.transformf], body_u: wp.array[wp.spatial_vectorf], @@ -538,7 +539,7 @@ def _apply_floating_base_transform( # Early return based on mask or absence of floating base wid = body_world_id[body_id] - if not world_mask[wid] or model_base_body_index[wid] < 0: + if (world_mask and not world_mask[wid]) or model_base_body_index[wid] < 0: return # Transform body pose @@ -569,7 +570,7 @@ def reset_time( model: ModelKamino, time: wp.array[wp.float32], steps: wp.array[wp.int32], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool] | None = None, ): wp.launch( _reset_time_of_select_worlds, @@ -590,7 +591,7 @@ def get_base_q_from_joint_q_and_body_q( joint_q: wp.array[wp.float32], body_q: wp.array[wp.transformf], base_q: wp.array[wp.transformf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool] | None = None, ): """ Infer the floating base pose from joint coordinates, if a base joint was set, or from body poses, @@ -601,7 +602,7 @@ def get_base_q_from_joint_q_and_body_q( joint_q: joint coordinates array. body_q: body poses array. base_q: array of per-world floating base pose, to set from joint_q/body_q as applicable. - world_mask: Per-world boolean mask, indicating in which worlds to perform the operation. + world_mask: Per-world boolean mask. If provided, indicates in which worlds to perform the operation. """ wp.launch( _get_base_q_from_joint_q_and_body_q, @@ -624,7 +625,7 @@ def get_base_u_from_joint_u_and_body_u( joint_u: wp.array[wp.float32], body_u: wp.array[wp.spatial_vectorf], base_u: wp.array[wp.spatial_vectorf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool] | None = None, ): """ Infer the floating base velocity from joint velocities, if a base joint was set, or from body velocities, @@ -635,7 +636,7 @@ def get_base_u_from_joint_u_and_body_u( joint_u: joint velocities array. body_u: body velocities array. base_u: array of per-world floating base velocity, to set from joint_u/body_u as applicable. - world_mask: Per-world boolean mask, indicating in which worlds to perform the operation. + world_mask: Per-world boolean mask. If provided, indicates in which worlds to perform the operation. """ wp.launch( _get_base_u_from_joint_u_and_body_u, @@ -657,7 +658,7 @@ def set_body_q( model: ModelKamino, body_q_in: wp.array[wp.transformf], body_q_out: wp.array[wp.transformf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool] | None = None, ): """ Set the body poses of select worlds to prescribed values. @@ -666,7 +667,7 @@ def set_body_q( model: Kamino model. body_q_in: prescribed body poses. body_q_out: body poses to overwrite with those in body_q_in, in active worlds. - world_mask: Per-world boolean mask, indicating in which worlds to perform the operation. + world_mask: Per-world boolean mask. If provided, indicates in which worlds to perform the operation. """ wp.launch( _set_body_q, @@ -682,7 +683,7 @@ def set_floating_base( base_u: wp.array[wp.spatial_vectorf] | None, body_q: wp.array[wp.transformf], body_u: wp.array[wp.spatial_vectorf], - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool] | None = None, relative_base_u: bool = False, ): """ @@ -697,7 +698,7 @@ def set_floating_base( If None, no additional velocity is composed to match the base velocity. body_q: body poses to update. body_u: body velocities to update. - world_mask: Per-world boolean mask, indicating in which worlds to perform the operation. + world_mask: Per-world boolean mask. If provided, indicates in which worlds to perform the operation. relative_base_u: Boolean indicating whether base_u should be interpreted as expressed relative to the new pose (after transforming so as to match base_q). """ @@ -754,7 +755,7 @@ def set_floating_base( def reset_joints_state_from_bodies_state( model: ModelKamino, state: StateKamino, - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool] | None = None, ): """ Reset joint-based components of the state given body poses and velocities, inferring consistent @@ -763,7 +764,7 @@ def reset_joints_state_from_bodies_state( Args: model: Kamino model. state: Kamino state. - world_mask: Per-world boolean mask, indicating in which worlds to perform the operation. + world_mask: Per-world boolean mask. If provided, indicates in which worlds to perform the operation. """ wp.launch( _reset_joints_state_from_bodies_state, @@ -796,7 +797,7 @@ def reset_joints_state_from_bodies_state( def reset_body_velocities( model: ModelKamino, state: StateKamino, - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool] | None = None, ): """ Reset body velocities in the state to zero. @@ -817,7 +818,7 @@ def reset_body_velocities( def reset_body_wrenches( model: ModelKamino, state: StateKamino, - world_mask: wp.array[wp.bool], + world_mask: wp.array[wp.bool] | None = None, ): """ Reset body wrenches in the state to zero. @@ -825,7 +826,7 @@ def reset_body_wrenches( Args: model: Kamino model. state: Kamino state. - world_mask: Per-world boolean mask, indicating in which worlds to perform the operation. + world_mask: Per-world boolean mask. If provided, indicates in which worlds to perform the operation. """ wp.launch( _reset_body_wrenches, diff --git a/newton/_src/solvers/kamino/_src/linalg/conjugate.py b/newton/_src/solvers/kamino/_src/linalg/conjugate.py index bd3278a5b6..def8031c48 100644 --- a/newton/_src/solvers/kamino/_src/linalg/conjugate.py +++ b/newton/_src/solvers/kamino/_src/linalg/conjugate.py @@ -338,7 +338,7 @@ def _run_capturable_loop( maxiter: wp.array[int], atol_sq: wp.array[Any], callback: Callable | None, - use_cuda_graph: bool, + use_graph: bool, use_graph_conditionals: bool = True, maxiter_host: int | None = None, loop_granularity: int = 1, @@ -384,7 +384,7 @@ def do_cycle_with_condition(): if callback_launch is not None: callback_launch.launch() - if use_cuda_graph and device.is_cuda and device.is_capturing: + if use_graph and device.is_capturing: if use_graph_conditionals: wp.capture_while(global_condition, do_cycle_with_condition) else: @@ -415,7 +415,7 @@ def less_than_op(i: wp.int32, threshold: wp.int32) -> wp.float32: def make_dot_kernel(tile_size: int, maxdim: int): num_tiles = (maxdim + tile_size - 1) // tile_size - @wp.kernel + @wp.kernel(module="unique", module_options={"enable_backward": False, "default_grid_stride": False}) def dot( a: wp.array2d[Any], b: wp.array2d[Any], @@ -541,7 +541,7 @@ class ConjugateSolver(Generic[ScalarType, IndexType]): maxiter: Maximum iterations per world. If None, defaults to 1.5 * maxdims. Mi: Operator applying the inverse preconditioner M^-1, such that Mi @ A has a smaller condition number than A. callback: Optional callback kernel invoked each iteration. - use_cuda_graph: Whether to use CUDA graph capture for the solve loop. + use_graph: Whether to use graph capture for the solve loop. loop_granularity: Number of iterations before termination criteria are checked. """ @@ -555,7 +555,7 @@ def __init__( maxiter: wp.array[wp.int32] | None = None, Mi: BatchedLinearOperator[ScalarType, IndexType] | None = None, callback: Callable | None = None, - use_cuda_graph: bool = True, + use_graph: bool = True, use_graph_conditionals: bool = True, loop_granularity: int = 1, ): @@ -593,7 +593,7 @@ def __init__( self.loop_granularity = loop_granularity self.callback = callback - self.use_cuda_graph = use_cuda_graph + self.use_graph = use_graph self.dot_tile_size = min(2048, 2 ** math.ceil(math.log(self.maxdims, 2))) self.tiled_dot_kernel = make_dot_kernel(self.dot_tile_size, self.maxdims) @@ -756,7 +756,7 @@ def solve( self.maxiter, self.atol_sq, self.callback, - self.use_cuda_graph, + self.use_graph, use_graph_conditionals=self.use_graph_conditionals, maxiter_host=self.maxiter_host, loop_granularity=min(self.loop_granularity, self.maxiter_host), @@ -890,7 +890,7 @@ def solve( self.maxiter, self.atol_sq, self.callback, - self.use_cuda_graph, + self.use_graph, use_graph_conditionals=self.use_graph_conditionals, maxiter_host=self.maxiter_host, loop_granularity=min(self.loop_granularity, self.maxiter_host), diff --git a/newton/_src/solvers/kamino/_src/linalg/factorize/_tile_builtins.py b/newton/_src/solvers/kamino/_src/linalg/factorize/_tile_builtins.py index 0d34720076..98179e8237 100644 --- a/newton/_src/solvers/kamino/_src/linalg/factorize/_tile_builtins.py +++ b/newton/_src/solvers/kamino/_src/linalg/factorize/_tile_builtins.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers # SPDX-License-Identifier: Apache-2.0 -"""Availability checks and native fallbacks for optional Warp tile builtins.""" +"""Availability checks and fallback for optional Warp tile builtins.""" import inspect import os @@ -26,157 +26,39 @@ def _has_warp_builtin(name: str) -> bool: return name in builtin_functions -def _has_native_tile_arg_support() -> bool: +def _has_native_left_transpose_update_support() -> bool: if not _tile_transpose_update_enabled(): return False try: from warp._src import codegen # noqa: PLC0415 - source = inspect.getsource(codegen.codegen_snippet) - except Exception: - return False - - return "template_params" in source and "is_tile(arg.type)" in source and "& {arg.emit()" in source - - -def _has_native_tile_access_helpers() -> bool: - if not _tile_transpose_update_enabled(): - return False - - try: + codegen_source = inspect.getsource(codegen.codegen_snippet) tile_header = Path(wp.__file__).resolve().parent / "native" / "tile.h" - source = tile_header.read_text(encoding="utf-8") + tile_source = tile_header.read_text(encoding="utf-8") except Exception: return False return ( - "tile_read(const tile_register_t" in source - and "tile_read(const tile_shared_t" in source - and "tile_add(tile_register_t" in source - and "tile_add(tile_shared_t" in source + "template_params" in codegen_source + and "is_tile(arg.type)" in codegen_source + and "& {arg.emit()" in codegen_source + and "tile_add(tile_register_t" in tile_source + and "tile_add(tile_shared_t" in tile_source ) -def _has_native_tile_update_support() -> bool: - return _has_native_tile_arg_support() and _has_native_tile_access_helpers() - - -def _copy_dense_2d_snippet(tile_name: str, layout_name: str, values_name: str, cols_name: str, storage: str) -> str: - if storage == "register": - return ( - f"{tile_name}.apply([&](int reg, auto c) {{ " - f"{values_name}[c[0] * {cols_name} + c[1]] = {tile_name}.data[reg]; }});" - ) - if storage == "generic": - return f"""for (int linear = WP_TILE_THREAD_IDX; linear < {layout_name}::Size; linear += WP_TILE_BLOCK_DIM) {{ - auto c = {layout_name}::coord_from_linear(linear); - int reg = linear / WP_TILE_BLOCK_DIM; - {values_name}[c[0] * {cols_name} + c[1]] = tile_read({tile_name}, reg, linear); -}}""" - raise ValueError(f"Unsupported tile storage specialization: {storage!r}") - - -def _update_output_snippet( - layout_name: str, - output_name: str, - rows_name: str, - cols_name: str, - k_name: str, - left_values_name: str, - right_values_name: str, - left_transposed: bool, - storage: str, -) -> str: - if left_transposed: - product = f"{left_values_name}[k * {rows_name} + c[0]] * {right_values_name}[k * {cols_name} + c[1]]" - else: - product = f"{left_values_name}[c[0] * {k_name} + k] * {right_values_name}[c[1] * {k_name} + k]" - - if storage == "shared": - write = f"""const T value = a * sum; - if constexpr ({layout_name}::Unique) - {output_name}.data(linear) += value; - else - wp::atomic_add(&{output_name}.data(linear), value);""" - elif storage == "register": - return f"""const T a = static_cast(alpha); -{output_name}.apply([&](int reg, auto c) {{ - T sum = T{{}}; - WP_PRAGMA_UNROLL - for (int k = 0; k < {k_name}; ++k) {{ - sum += {product}; - }} - {output_name}.data[reg] += a * sum; -}}); -WP_TILE_SYNC();""" - elif storage == "generic": - write = f"""int reg = linear / WP_TILE_BLOCK_DIM; - tile_add({output_name}, reg, linear, a * sum);""" - else: - raise ValueError(f"Unsupported tile storage specialization: {storage!r}") - - return f"""const T a = static_cast(alpha); -for (int linear = WP_TILE_THREAD_IDX; linear < {layout_name}::Size; linear += WP_TILE_BLOCK_DIM) {{ - auto c = {layout_name}::coord_from_linear(linear); - T sum = T{{}}; - WP_PRAGMA_UNROLL - for (int k = 0; k < {k_name}; ++k) {{ - sum += {product}; - }} - {write} -}} -WP_TILE_SYNC();""" - - -def _make_tile_matmul_transpose_update_snippet(out_storage: str, input_storage: str) -> str: - copy_left = _copy_dense_2d_snippet("left", "LeftLayout", "left_values", "K", input_storage) - copy_right = _copy_dense_2d_snippet("right", "RightLayout", "right_values", "K", input_storage) - update_out = _update_output_snippet( - "OutLayout", "out", "Rows", "Cols", "K", "left_values", "right_values", False, out_storage - ) - return f"""using OutTile = tile_out; -using LeftTile = tile_left; -using RightTile = tile_right; -using T = typename OutTile::Type; -using OutLayout = typename OutTile::Layout; -using LeftLayout = typename LeftTile::Layout; -using RightLayout = typename RightTile::Layout; - -static_assert(OutLayout::Shape::N == 2, "out must be 2D"); -static_assert(LeftLayout::Shape::N == 2, "left must be 2D"); -static_assert(RightLayout::Shape::N == 2, "right must be 2D"); -static_assert(LeftLayout::Shape::dim(0) == OutLayout::Shape::dim(0), "left rows must match out rows"); -static_assert(RightLayout::Shape::dim(0) == OutLayout::Shape::dim(1), "right rows must match out cols"); -static_assert(LeftLayout::Shape::dim(1) == RightLayout::Shape::dim(1), "left/right cols must match"); - -constexpr int Rows = OutLayout::Shape::dim(0); -constexpr int Cols = OutLayout::Shape::dim(1); -constexpr int K = LeftLayout::Shape::dim(1); - -#if defined(__CUDA_ARCH__) -__shared__ T left_values[Rows * K]; -__shared__ T right_values[Cols * K]; -#else -T left_values[Rows * K]; -T right_values[Cols * K]; -#endif - -{copy_left} -{copy_right} -WP_TILE_SYNC(); - -{update_out} -""" +HAS_TILE_MATMUL_TRANSPOSE_UPDATE = _has_warp_builtin("tile_matmul_transpose_update") +HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE = _has_warp_builtin("tile_matmul_left_transpose_update") +HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE = ( + not HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE and _has_native_left_transpose_update_support() +) -def _make_tile_matmul_left_transpose_update_snippet(out_storage: str, input_storage: str) -> str: - copy_left = _copy_dense_2d_snippet("left", "LeftLayout", "left_values", "Rows", input_storage) - copy_right = _copy_dense_2d_snippet("right", "RightLayout", "right_values", "Cols", input_storage) - update_out = _update_output_snippet( - "OutLayout", "out", "Rows", "Cols", "K", "left_values", "right_values", True, out_storage - ) - return f"""using OutTile = tile_out; +@cache +def make_tile_matmul_left_transpose_update_func(block_size: int): + """Create ``out += alpha * transpose(left) @ right`` for the LLT solve.""" + snippet = """using OutTile = tile_out; using LeftTile = tile_left; using RightTile = tile_right; using T = typename OutTile::Type; @@ -203,46 +85,22 @@ def _make_tile_matmul_left_transpose_update_snippet(out_storage: str, input_stor T right_values[K * Cols]; #endif -{copy_left} -{copy_right} +left.apply([&](int reg, auto c) { left_values[c[0] * Rows + c[1]] = left.data[reg]; }); +right.apply([&](int reg, auto c) { right_values[c[0] * Cols + c[1]] = right.data[reg]; }); WP_TILE_SYNC(); -{update_out} -""" - - -HAS_TILE_MATMUL_TRANSPOSE_UPDATE = _has_warp_builtin("tile_matmul_transpose_update") -HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE = _has_warp_builtin("tile_matmul_left_transpose_update") -HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE = not HAS_TILE_MATMUL_TRANSPOSE_UPDATE and _has_native_tile_update_support() -HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE = ( - not HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE and _has_native_tile_update_support() -) - - -@cache -def make_tile_matmul_transpose_update_func( - block_size: int, out_storage: str = "shared", input_storage: str = "register" -): - """Create ``out += alpha * left @ transpose(right)`` as a native tile function.""" - snippet = _make_tile_matmul_transpose_update_snippet(out_storage, input_storage) - - @wp.func_native(snippet) - def tile_matmul_transpose_update( - out: wp.tile[float, block_size, block_size], - left: wp.tile[float, block_size, block_size], - right: wp.tile[float, block_size, block_size], - alpha: float, - ): ... - - return tile_matmul_transpose_update - - -@cache -def make_tile_matmul_left_transpose_update_func( - block_size: int, out_storage: str = "generic", input_storage: str = "register" -): - """Create ``out += alpha * transpose(left) @ right`` as a native tile function.""" - snippet = _make_tile_matmul_left_transpose_update_snippet(out_storage, input_storage) +const T a = static_cast(alpha); +for (int linear = WP_TILE_THREAD_IDX; linear < OutLayout::Size; linear += WP_TILE_BLOCK_DIM) { + auto c = OutLayout::coord_from_linear(linear); + T sum = T{}; + WP_PRAGMA_UNROLL + for (int k = 0; k < K; ++k) { + sum += left_values[k * Rows + c[0]] * right_values[k * Cols + c[1]]; + } + int reg = linear / WP_TILE_BLOCK_DIM; + tile_add(out, reg, linear, a * sum); +} +WP_TILE_SYNC();""" @wp.func_native(snippet) def tile_matmul_left_transpose_update( diff --git a/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked.py b/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked.py index aa7494da4b..b0f2dcf371 100644 --- a/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked.py +++ b/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked.py @@ -10,11 +10,9 @@ from ._tile_builtins import ( HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE, - HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE, HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE, HAS_TILE_MATMUL_TRANSPOSE_UPDATE, make_tile_matmul_left_transpose_update_func, - make_tile_matmul_transpose_update_func, ) ### @@ -166,10 +164,6 @@ def llt_blocked_factorize_kernel( L_block = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, j)) if wp.static(HAS_TILE_MATMUL_TRANSPOSE_UPDATE): wp.tile_matmul_transpose_update(A_kk_tile, L_block, L_block, alpha=-1.0) - elif wp.static(HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE): - wp.static(make_tile_matmul_transpose_update_func(block_size, "shared", "register"))( - A_kk_tile, L_block, L_block, -1.0 - ) else: L_block_T = wp.tile_transpose(L_block) wp.tile_matmul(L_block, L_block_T, A_kk_tile, alpha=-1.0) @@ -206,10 +200,6 @@ def llt_blocked_factorize_kernel( L_2_tile = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, j)) if wp.static(HAS_TILE_MATMUL_TRANSPOSE_UPDATE): wp.tile_matmul_transpose_update(A_ik_tile, L_tile, L_2_tile, alpha=-1.0) - elif wp.static(HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE): - wp.static(make_tile_matmul_transpose_update_func(block_size, "shared", "register"))( - A_ik_tile, L_tile, L_2_tile, -1.0 - ) else: L_T_tile = wp.tile_transpose(L_2_tile) wp.tile_matmul(L_tile, L_T_tile, A_ik_tile, alpha=-1.0) @@ -306,7 +296,7 @@ def llt_blocked_solve_kernel( if wp.static(HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE): wp.tile_matmul_left_transpose_update(rhs_tile, L_tile, x_tile, alpha=-1.0) elif wp.static(HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE): - wp.static(make_tile_matmul_left_transpose_update_func(block_size, "generic", "register"))( + wp.static(make_tile_matmul_left_transpose_update_func(block_size))( rhs_tile, L_tile, x_tile, -1.0 ) else: @@ -395,7 +385,7 @@ def llt_blocked_solve_inplace_kernel( if wp.static(HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE): wp.tile_matmul_left_transpose_update(rhs_tile, L_tile, x_tile, alpha=-1.0) elif wp.static(HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE): - wp.static(make_tile_matmul_left_transpose_update_func(block_size, "generic", "register"))( + wp.static(make_tile_matmul_left_transpose_update_func(block_size))( rhs_tile, L_tile, x_tile, -1.0 ) else: diff --git a/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_rcm.py b/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_rcm.py index 8137de1fc9..7011b1dfea 100644 --- a/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_rcm.py +++ b/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_rcm.py @@ -36,11 +36,9 @@ from ._tile_builtins import ( HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE, - HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE, HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE, HAS_TILE_MATMUL_TRANSPOSE_UPDATE, make_tile_matmul_left_transpose_update_func, - make_tile_matmul_transpose_update_func, ) ### @@ -56,6 +54,7 @@ "llt_blocked_rcm_symbolic_fill_in", "make_llt_blocked_rcm_factorize_kernel", "make_llt_blocked_rcm_fused_permute_and_tp_kernel", + "make_llt_blocked_rcm_parallel_factorize_kernels", "make_llt_blocked_rcm_permute_vector_kernel", "make_llt_blocked_rcm_solve_inplace_kernel", "make_llt_blocked_rcm_solve_kernel", @@ -132,7 +131,8 @@ def make_llt_blocked_rcm_fused_permute_and_tp_kernel(block_size: int, max_dim: i """Fused kernel: builds ``inv_P``, permutes ``A -> A_hat``, and reduces ``|A_hat|`` into the tile pattern in a single launch. - Launch dims: ``(num_blocks, max_dim, max_dim)``. Each thread ``(b, r, c)``: + Launch dims: ``(num_blocks, max_dim * (max_dim + 1) // 2)``. Each thread + processes one element of the lower triangle: 1. If ``c == 0``: writes ``inv_P[P[r]] = r`` for block ``b``. 2. Computes ``v = A[P[r], P[c]]`` and writes it into ``A_hat[r, c]``. @@ -162,10 +162,23 @@ def fused_permute_and_tp_kernel( inv_P: wp.array[wp.int32], tile_pattern: wp.array[wp.int32], ): - b, r, c = wp.tid() + b, triangular_index = wp.tid() n_i = dim[b] - if r >= n_i or c >= n_i: + triangular_size = n_i * (n_i + 1) // 2 + if triangular_index >= triangular_size: return + + # Dense systems larger than 23169 can overflow int32 in 8 * triangular_index; + # systems at that scale should use the sparse factorization path. + r = int((wp.sqrt(float(8 * triangular_index + 1)) - float(1)) * float(0.5)) + row_start = r * (r + 1) // 2 + if row_start > triangular_index: + r -= 1 + row_start = r * (r + 1) // 2 + elif (r + 1) * (r + 2) // 2 <= triangular_index: + r += 1 + row_start = r * (r + 1) // 2 + c = triangular_index - row_start mat_off = mio[b] vec_off = vio[b] tp_off = tpo[b] @@ -320,10 +333,6 @@ def llt_blocked_rcm_factorize_kernel( L_block = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, j)) if wp.static(HAS_TILE_MATMUL_TRANSPOSE_UPDATE): wp.tile_matmul_transpose_update(A_kk_tile, L_block, L_block, alpha=-1.0) - elif wp.static(HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE): - wp.static(make_tile_matmul_transpose_update_func(block_size, "shared", "register"))( - A_kk_tile, L_block, L_block, -1.0 - ) else: L_block_T = wp.tile_transpose(L_block) wp.tile_matmul(L_block, L_block_T, A_kk_tile, alpha=-1.0) @@ -365,10 +374,6 @@ def llt_blocked_rcm_factorize_kernel( L_2_tile = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, j)) if wp.static(HAS_TILE_MATMUL_TRANSPOSE_UPDATE): wp.tile_matmul_transpose_update(A_ik_tile, L_tile, L_2_tile, alpha=-1.0) - elif wp.static(HAS_NATIVE_TILE_MATMUL_TRANSPOSE_UPDATE): - wp.static(make_tile_matmul_transpose_update_func(block_size, "shared", "register"))( - A_ik_tile, L_tile, L_2_tile, -1.0 - ) else: L_T_tile = wp.tile_transpose(L_2_tile) wp.tile_matmul(L_tile, L_T_tile, A_ik_tile, alpha=-1.0) @@ -381,12 +386,135 @@ def llt_blocked_rcm_factorize_kernel( return llt_blocked_rcm_factorize_kernel +@cache +def make_llt_blocked_rcm_parallel_factorize_kernels(block_size: int): + """Create panel-parallel blocked Cholesky kernels. + + Each diagonal tile remains sequential, but the off-diagonal tiles in a + panel are solved by independent CUDA blocks. This exposes parallelism for + a single large matrix while preserving the same factor and tile mask. + """ + + @wp.kernel(enable_backward=False) + def factorize_diagonal_kernel( + tile_k: int, + dim: wp.array[wp.int32], + mio: wp.array[wp.int32], + tpo: wp.array[wp.int32], + A: wp.array[wp.float32], + tile_pattern: wp.array[wp.int32], + L: wp.array[wp.float32], + ): + bid, tid_block = wp.tid() + block_dim = wp.block_dim() + n = dim[bid] + k = tile_k * block_size + if k >= n: + return + + mat_offset = mio[bid] + pattern_offset = tpo[bid] + A_i = wp.array(ptr=get_float32_array_offset_ptr(A, mat_offset), shape=(n, n), dtype=wp.float32) + L_i = wp.array(ptr=get_float32_array_offset_ptr(L, mat_offset), shape=(n, n), dtype=wp.float32) + n_tiles = (n + block_size - 1) // block_size + TP_i = wp.array( + ptr=get_int32_array_offset_ptr(tile_pattern, pattern_offset), + shape=(n_tiles, n_tiles), + dtype=wp.int32, + ) + + diagonal = wp.tile_load(A_i, shape=(block_size, block_size), offset=(k, k), storage="shared") + if k + block_size > n: + for q in range((block_size * block_size + block_dim - 1) // block_dim): + index = (tid_block + q * block_dim) % (block_size * block_size) + row = index // block_size + col = index % block_size + # Preserve a collective full-tile write before the next Tile operation. + value = diagonal[row, col] + if k + row >= n or k + col >= n: + value = wp.where(row == col, wp.float32(1), wp.float32(0)) + diagonal[row, col] = value + + for tile_j in range(tile_k): + if TP_i[tile_k, tile_j] == int(0): + continue + j = tile_j * block_size + previous = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, j)) + wp.tile_matmul(previous, wp.tile_transpose(previous), diagonal, alpha=-1.0) + + wp.tile_cholesky_inplace(diagonal) + wp.tile_store(L_i, diagonal, offset=(k, k)) + + @wp.kernel(enable_backward=False) + def factorize_panel_kernel( + tile_k: int, + dim: wp.array[wp.int32], + mio: wp.array[wp.int32], + tpo: wp.array[wp.int32], + A: wp.array[wp.float32], + tile_pattern: wp.array[wp.int32], + L: wp.array[wp.float32], + ): + bid, panel_tile_i, tid_block = wp.tid() + tile_i = panel_tile_i + tile_k + 1 + block_dim = wp.block_dim() + n = dim[bid] + n_tiles = (n + block_size - 1) // block_size + if tile_i >= n_tiles: + return + + mat_offset = mio[bid] + pattern_offset = tpo[bid] + A_i = wp.array(ptr=get_float32_array_offset_ptr(A, mat_offset), shape=(n, n), dtype=wp.float32) + L_i = wp.array(ptr=get_float32_array_offset_ptr(L, mat_offset), shape=(n, n), dtype=wp.float32) + TP_i = wp.array( + ptr=get_int32_array_offset_ptr(tile_pattern, pattern_offset), + shape=(n_tiles, n_tiles), + dtype=wp.int32, + ) + if TP_i[tile_i, tile_k] == int(0): + return + + i = tile_i * block_size + k = tile_k * block_size + panel = wp.tile_load(A_i, shape=(block_size, block_size), offset=(i, k), storage="shared") + diagonal = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, k), storage="shared") + if i + block_size > n or k + block_size > n: + for q in range((block_size * block_size + block_dim - 1) // block_dim): + index = (tid_block + q * block_dim) % (block_size * block_size) + row = index // block_size + col = index % block_size + # Preserve collective full-tile writes before the next Tile operations. + panel_value = panel[row, col] + if i + row >= n or k + col >= n: + panel_value = wp.where(i + row == k + col, wp.float32(1), wp.float32(0)) + panel[row, col] = panel_value + diagonal_value = diagonal[row, col] + if k + row >= n or k + col >= n: + diagonal_value = wp.where(row == col, wp.float32(1), wp.float32(0)) + diagonal[row, col] = diagonal_value + + for tile_j in range(tile_k): + if TP_i[tile_i, tile_j] == int(0) or TP_i[tile_k, tile_j] == int(0): + continue + j = tile_j * block_size + left = wp.tile_load(L_i, shape=(block_size, block_size), offset=(i, j)) + right = wp.tile_load(L_i, shape=(block_size, block_size), offset=(k, j)) + wp.tile_matmul(left, wp.tile_transpose(right), panel, alpha=-1.0) + + transposed = wp.tile_transpose(panel) + wp.tile_lower_solve_inplace(diagonal, transposed) + wp.tile_store(L_i, wp.tile_transpose(transposed), offset=(i, k)) + + return factorize_diagonal_kernel, factorize_panel_kernel + + @cache def make_llt_blocked_rcm_solve_kernel(block_size: int): """RCM solve with tile skipping and fused output un-permutation. - The RHS is already in permuted coordinates. The solve writes ``x_hat`` in - permuted coordinates for backward-substitution dependencies and scatters + The solve gathers the RHS into permuted coordinates, writes ``x_hat`` in + permuted coordinates for backward-substitution dependencies, and scatters each solved tile directly to the original-coordinate output ``x``. """ @@ -436,7 +564,15 @@ def llt_blocked_rcm_solve_kernel( # Forward substitution: solve L y = b. for i in range(0, n_i_padded, block_size): tile_i = i // block_size - rhs_tile = wp.tile_load(b_i, shape=(block_size, 1), offset=(i, 0)) + rhs_tile = wp.tile_zeros(shape=(block_size, 1), dtype=wp.float32, storage="shared") + num_row_iterations = (block_size + num_threads_per_block - 1) // num_threads_per_block + for ii in range(num_row_iterations): + row = tid_block + ii * num_threads_per_block + active = row < block_size and i + row < n_i + value = wp.float32(0.0) + if active: + value = b_i[P_i[i + row], 0] + wp.tile_scatter_masked(rhs_tile, row, 0, value, active) L_diag = wp.tile_load(L_i, shape=(block_size, block_size), offset=(i, i)) if i > 0: for j in range(0, i, block_size): @@ -479,7 +615,7 @@ def llt_blocked_rcm_solve_kernel( if wp.static(HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE): wp.tile_matmul_left_transpose_update(rhs_tile, L_tile, x_tile, alpha=-1.0) elif wp.static(HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE): - wp.static(make_tile_matmul_left_transpose_update_func(block_size, "generic", "register"))( + wp.static(make_tile_matmul_left_transpose_update_func(block_size))( rhs_tile, L_tile, x_tile, -1.0 ) else: @@ -587,7 +723,7 @@ def llt_blocked_rcm_solve_inplace_kernel( if wp.static(HAS_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE): wp.tile_matmul_left_transpose_update(rhs_tile, L_tile, x_tile, alpha=-1.0) elif wp.static(HAS_NATIVE_TILE_MATMUL_LEFT_TRANSPOSE_UPDATE): - wp.static(make_tile_matmul_left_transpose_update_func(block_size, "generic", "register"))( + wp.static(make_tile_matmul_left_transpose_update_func(block_size))( rhs_tile, L_tile, x_tile, -1.0 ) else: @@ -648,7 +784,7 @@ def llt_blocked_rcm_fused_permute_and_tp( """ wp.launch( kernel=kernel, - dim=(num_blocks, max_dim, max_dim), + dim=(num_blocks, max_dim * (max_dim + 1) // 2), inputs=[dim, mio, vio, tpo, float(tol), P, A, A_hat, inv_P, tile_pattern], device=device, ) @@ -694,6 +830,40 @@ def llt_blocked_rcm_factorize( ) +def llt_blocked_rcm_factorize_parallel( + kernels, + dim: wp.array[wp.int32], + mio: wp.array[wp.int32], + tpo: wp.array[wp.int32], + A: wp.array[wp.float32], + tile_pattern: wp.array[wp.int32], + L: wp.array[wp.float32], + num_blocks: int, + max_tiles: int, + block_dim: int = 128, + device: wp.DeviceLike = None, +): + """Launch the panel-parallel semi-sparse blocked Cholesky factorization.""" + diagonal_kernel, panel_kernel = kernels + for tile_k in range(max_tiles): + wp.launch_tiled( + kernel=diagonal_kernel, + dim=num_blocks, + inputs=[tile_k, dim, mio, tpo, A, tile_pattern, L], + block_dim=block_dim, + device=device, + ) + panel_tiles = max_tiles - tile_k - 1 + if panel_tiles > 0: + wp.launch_tiled( + kernel=panel_kernel, + dim=(num_blocks, panel_tiles), + inputs=[tile_k, dim, mio, tpo, A, tile_pattern, L], + block_dim=block_dim, + device=device, + ) + + def llt_blocked_rcm_solve( kernel, dim: wp.array[wp.int32], diff --git a/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_rcm_solver.py b/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_rcm_solver.py index 832b66af0d..19c11d00b2 100644 --- a/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_rcm_solver.py +++ b/newton/_src/solvers/kamino/_src/linalg/factorize/llt_blocked_rcm_solver.py @@ -35,6 +35,7 @@ from . import rcm_batch as _rcm_batch from .llt_blocked_rcm import ( llt_blocked_rcm_factorize, + llt_blocked_rcm_factorize_parallel, llt_blocked_rcm_fused_permute_and_tp, llt_blocked_rcm_permute_vector, llt_blocked_rcm_solve, @@ -42,6 +43,7 @@ llt_blocked_rcm_symbolic_fill_in, make_llt_blocked_rcm_factorize_kernel, make_llt_blocked_rcm_fused_permute_and_tp_kernel, + make_llt_blocked_rcm_parallel_factorize_kernels, make_llt_blocked_rcm_permute_vector_kernel, make_llt_blocked_rcm_solve_inplace_kernel, make_llt_blocked_rcm_solve_kernel, @@ -105,8 +107,10 @@ def __init__( # Threshold below which ``|A[i,j]|`` is treated as a non-edge by the # RCM adjacency scan and by the tile-pattern builder. reorder_tol: float = 0.0, - # Cap on BFS steps per block. None => auto (``2*ceil(sqrt(n)) + 4``). + # Optional approximate traversal cap. None completes every component. rcm_max_bfs_iters: int | None = None, + reuse_permutation: bool = True, + parallel_factorization: bool = False, dtype: FloatType = wp.float32, device: wp.DeviceLike | None = None, **kwargs: dict[str, Any], @@ -120,7 +124,13 @@ def __init__( reorder_tol: threshold below which an off-diagonal entry is treated as a non-edge by the RCM adjacency scan and by the tile-pattern builder. - rcm_max_bfs_iters: BFS depth cap for the batched RCM pass. + rcm_max_bfs_iters: optional BFS step cap. By default every connected + component is traversed completely. + reuse_permutation: whether to compute RCM once and reuse that + permutation for later numeric factorizations. The numeric tile + pattern is still rebuilt each time. Defaults to ``True``. + parallel_factorization: whether to solve off-diagonal tiles of + each Cholesky panel in parallel. Defaults to ``False``. """ # The underlying kernels (factorize / solve / permute / tile-pattern) # are hard-coded to wp.float32, so reject any other dtype up front @@ -133,7 +143,6 @@ def __init__( self._y: wp.array[dtype] | None = None # Reordering + semi-sparse state self._A_hat: wp.array[dtype] | None = None - self._b_hat: wp.array[dtype] | None = None self._x_hat: wp.array[dtype] | None = None self._P: wp.array[wp.int32] | None = None self._inv_P: wp.array[wp.int32] | None = None @@ -159,9 +168,12 @@ def __init__( # Reordering options self._reorder_tol: float = reorder_tol self._rcm_max_bfs_iters = rcm_max_bfs_iters + self._reuse_permutation = reuse_permutation + self._parallel_factorization = parallel_factorization # Build kernels (cached by block_size / max_dim at allocate time). self._factorize_kernel = make_llt_blocked_rcm_factorize_kernel(block_size) + self._parallel_factorize_kernels = make_llt_blocked_rcm_parallel_factorize_kernels(block_size) self._solve_kernel = make_llt_blocked_rcm_solve_kernel(block_size) self._solve_inplace_kernel = make_llt_blocked_rcm_solve_inplace_kernel(block_size) # Auxiliary kernels resolved in _allocate_impl once we know max_dim. @@ -256,7 +268,6 @@ def _allocate_impl(self, A: DenseLinearOperatorData[wp.float32, wp.int32], **kwa # Reordering scratch. self._A_hat = wp.zeros(shape=(info.total_mat_size,), dtype=self._dtype) - self._b_hat = wp.zeros(shape=(info.total_vec_size,), dtype=self._dtype) self._x_hat = wp.zeros(shape=(info.total_vec_size,), dtype=self._dtype) # Permutations (indexed by vio, length dim per block). @@ -288,9 +299,10 @@ def _reset_impl(self) -> None: self._L.zero_() self._y.zero_() self._A_hat.zero_() - self._b_hat.zero_() self._x_hat.zero_() self._P.zero_() + self._rcm_scratch["permutation_valid"].zero_() + self._rcm_scratch["permutation_dim"].zero_() self._inv_P.zero_() self._tile_pattern.zero_() self._has_factors = False @@ -319,6 +331,7 @@ def _ensure_reorder_launches_bound(self, A: wp.array[Any]) -> None: tol=self._reorder_tol, max_bfs_iters=self._rcm_max_bfs_iters, use_cuda_graph=False, + reuse_permutation=self._reuse_permutation, device=self._device, ) self._reorder_attached_to = A @@ -331,9 +344,9 @@ def _factorize_impl(self, A: wp.array[Any]) -> None: # Bind / rebind views to the current A buffer. self._ensure_reorder_launches_bound(A) - # 1. Compute per-block P via the batched RCM callback. The callback - # is a set of recorded Warp launches and is safe to replay under - # CUDA graph capture initiated by the caller. + # Compute per-block P via the batched RCM callback. The callback is a + # set of recorded Warp launches and is safe to replay under CUDA graph + # capture initiated by the caller. self._reorder_callback() # 2. Fused: build inv_P, permute A -> A_hat, and reduce |A_hat| into the @@ -370,18 +383,33 @@ def _factorize_impl(self, A: wp.array[Any]) -> None: ) # 4. Numeric factorization with tile-pattern skips. - llt_blocked_rcm_factorize( - kernel=self._factorize_kernel, - dim=info.dim, - mio=info.mio, - tpo=self._tpo, - A=self._A_hat, - tile_pattern=self._tile_pattern, - L=self._L, - num_blocks=num_blocks, - block_dim=self._factorize_block_dim, - device=self._device, - ) + if self._parallel_factorization: + llt_blocked_rcm_factorize_parallel( + kernels=self._parallel_factorize_kernels, + dim=info.dim, + mio=info.mio, + tpo=self._tpo, + A=self._A_hat, + tile_pattern=self._tile_pattern, + L=self._L, + num_blocks=num_blocks, + max_tiles=(self._max_dim + self._block_size - 1) // self._block_size, + block_dim=self._factorize_block_dim, + device=self._device, + ) + else: + llt_blocked_rcm_factorize( + kernel=self._factorize_kernel, + dim=info.dim, + mio=info.mio, + tpo=self._tpo, + A=self._A_hat, + tile_pattern=self._tile_pattern, + L=self._L, + num_blocks=num_blocks, + block_dim=self._factorize_block_dim, + device=self._device, + ) @override def _reconstruct_impl(self, A: wp.array[Any]) -> None: @@ -392,20 +420,7 @@ def _solve_impl(self, b: wp.array[Any], x: wp.array[Any]) -> None: info = self._operator.info num_blocks = info.num_blocks - # Permute b -> b_hat. - llt_blocked_rcm_permute_vector( - kernel=self._permute_vector_kernel, - dim=info.dim, - vio=info.vio, - P=self._P, - src=b, - dst=self._b_hat, - num_blocks=num_blocks, - max_dim=self._max_dim, - device=self._device, - ) - - # Solve L L^T x_hat = b_hat and scatter x_hat -> x. + # Solve L L^T x_hat = P b and scatter x_hat -> x. llt_blocked_rcm_solve( kernel=self._solve_kernel, dim=info.dim, @@ -415,7 +430,7 @@ def _solve_impl(self, b: wp.array[Any], x: wp.array[Any]) -> None: P=self._P, L=self._L, tile_pattern=self._tile_pattern, - b=self._b_hat, + b=b, y=self._y, x_hat=self._x_hat, x=x, diff --git a/newton/_src/solvers/kamino/_src/linalg/factorize/rcm_batch.py b/newton/_src/solvers/kamino/_src/linalg/factorize/rcm_batch.py index 00e4f542a0..b866b71a51 100644 --- a/newton/_src/solvers/kamino/_src/linalg/factorize/rcm_batch.py +++ b/newton/_src/solvers/kamino/_src/linalg/factorize/rcm_batch.py @@ -14,11 +14,9 @@ in ``B * max_bfs_iters``. At small problem sizes (e.g. ``n = 256``, ``B = 8``) the resulting hundreds of launches dominate wall time over the actual compute. -The CUDA/float32 fast path runs each graph block inside one tiled Warp -kernel, using shared memory for the per-vertex RCM state and an in-kernel BFS -loop. Larger or non-CUDA cases fall back to the staged batched path, which -amortizes launch overhead by making each RCM stage a single launch that covers -**all** blocks. +The CUDA complete-traversal path keeps the full BFS inside one persistent +kernel per matrix block. Other devices use a staged fallback with a fixed +``max_dim`` upper bound whose completed iterations are no-ops. Layout assumptions ------------------ @@ -54,13 +52,20 @@ launch() # one zero-arg callable; CUDA-graph capturable """ -import math from collections.abc import Callable from functools import cache import warp as wp +@wp.func_native(""" +#if defined(__CUDA_ARCH__) +__syncthreads(); +#endif +""") +def _sync_threads(): ... + + def create_cuda_graph_callback(callback: Callable[[], None], device=None, stream=None) -> Callable[[], None]: """Capture ``callback`` into a CUDA graph and return a zero-arg replay fn.""" with wp.ScopedCapture(device=device, stream=stream) as capture: @@ -87,20 +92,18 @@ def allocate_rcm_batch_scratch(total_vec: int, num_blocks: int, device) -> dict: - Per-vertex arrays (``degree``, ``level``, ``order_buf``) are sized by the union of all block vector offsets (``total_vec = sum(dims)``). Each block's slice is ``[vio[b] : vio[b]+dims[b])``. - - Per-block arrays (``head``, ``root``) are sized ``(num_blocks,)`` - and are indexed by block id ``b``. - - The BFS "current level" scalar is *not* allocated here: it is a - host-side loop counter baked into each pre-recorded ``bfs_step`` - launch, so there is no device-side counter and no intra-launch race - hazard. + - Per-block arrays are sized ``(num_blocks,)`` and indexed by block. """ return { "degree": wp.empty(total_vec, dtype=wp.int32, device=device), "level": wp.empty(total_vec, dtype=wp.int32, device=device), "order_buf": wp.empty(total_vec, dtype=wp.int32, device=device), "head": wp.empty(num_blocks, dtype=wp.int32, device=device), - "root": wp.empty(num_blocks, dtype=wp.int32, device=device), + "current_level": wp.empty(num_blocks, dtype=wp.int32, device=device), + "discovered": wp.empty(num_blocks, dtype=wp.int32, device=device), + "reorder_active": wp.empty(num_blocks, dtype=wp.int32, device=device), + "permutation_valid": wp.zeros(num_blocks, dtype=wp.int32, device=device), + "permutation_dim": wp.zeros(num_blocks, dtype=wp.int32, device=device), } @@ -118,6 +121,19 @@ def _make_rcm_batch_kernels(dtype): module = wp.get_module(module_name) module.options.update({"enable_backward": False, "default_grid_stride": False}) + @wp.kernel(module=module) + def prepare_reorder_kernel( + reuse_permutation: bool, + dims: wp.array[wp.int32], + permutation_valid: wp.array[wp.int32], + permutation_dim: wp.array[wp.int32], + reorder_active: wp.array[wp.int32], + ): + """Mark blocks whose permutation must be computed.""" + b = wp.tid() + cached = permutation_valid[b] != int(0) and permutation_dim[b] == dims[b] + reorder_active[b] = wp.where(reuse_permutation and cached, int(0), int(1)) + @wp.kernel(module=module) def init_and_degree_kernel( num_blocks: int, @@ -129,7 +145,9 @@ def init_and_degree_kernel( degree: wp.array[wp.int32], # type: ignore[valid-type] level: wp.array[wp.int32], # type: ignore[valid-type] head: wp.array[wp.int32], # type: ignore[valid-type] - root: wp.array[wp.int32], # type: ignore[valid-type] + current_level: wp.array[wp.int32], + discovered: wp.array[wp.int32], + reorder_active: wp.array[wp.int32], ): """Launch dims: ``(num_blocks, max_dim)``. @@ -140,6 +158,8 @@ def init_and_degree_kernel( b, i = wp.tid() if b >= num_blocks: return + if reorder_active[b] == int(0): + return n_b = dims[b] if i >= n_b: return @@ -164,7 +184,8 @@ def init_and_degree_kernel( # Per-block scalars: one thread per block sets them. if i == 0: head[b] = int(0) - root[b] = int(0) + current_level[b] = int(0) + discovered[b] = int(0) @wp.kernel(module=module) def select_and_seed_kernel( @@ -175,21 +196,14 @@ def select_and_seed_kernel( level: wp.array[wp.int32], # type: ignore[valid-type] order_buf: wp.array[wp.int32], # type: ignore[valid-type] head: wp.array[wp.int32], # type: ignore[valid-type] - root: wp.array[wp.int32], # type: ignore[valid-type] + reorder_active: wp.array[wp.int32], ): - """Launch dims: ``(num_blocks,)``. Fused root-selection + BFS seed. - - One thread per block does a serialized argmin over that block's - ``degree`` slice to pick a minimum-degree root, stores it into - ``root[b]``, then seeds the BFS frontier for that block by writing - ``level[vb + r] = 0`` and appending ``r`` to the block's - ``order_buf`` segment (head advances by one). Fine for the intended - ``n <= ~1000`` regime; merging the two removes one kernel launch - at the start of every reorder call. - """ + """Select a minimum-degree root and seed each active block.""" b = wp.tid() if b >= num_blocks: return + if reorder_active[b] == int(0): + return n_b = dims[b] vb = vio[b] best_deg = int(2147483647) @@ -199,7 +213,6 @@ def select_and_seed_kernel( if d < best_deg: best_deg = d best_idx = i - root[b] = best_idx level[vb + best_idx] = int(0) # Atomically claim the first slot; at kernel entry head[b] is 0 and # only this thread touches it for block ``b``, so the atomic is @@ -210,7 +223,6 @@ def select_and_seed_kernel( @wp.kernel(module=module) def bfs_step_kernel( num_blocks: int, - cur: int, tol: dtype, # type: ignore[valid-type] A: wp.array[dtype], # type: ignore[valid-type] dims: wp.array[wp.int32], # type: ignore[valid-type] @@ -219,32 +231,27 @@ def bfs_step_kernel( level: wp.array[wp.int32], # type: ignore[valid-type] order_buf: wp.array[wp.int32], # type: ignore[valid-type] head: wp.array[wp.int32], # type: ignore[valid-type] + current_level: wp.array[wp.int32], + discovered: wp.array[wp.int32], + reorder_active: wp.array[wp.int32], ): """Launch dims: ``(num_blocks, max_dim)``. One BFS expansion step. - The "current BFS level" ``cur`` is passed as a **scalar kernel - argument** rather than being stored in device memory. The host - pre-records one launch per iteration with a distinct integer - ``cur`` baked into each, so every thread in a given launch - observes the same value. This removes the previous device-side - ``iter_counter`` scratch array, the per-step 1-thread increment - launch, and any possibility of an intra-launch race on the - counter. - - The ``alive`` / ``discovered`` arrays are dropped entirely: when a - block saturates, no thread in it has ``level == cur`` so the kernel - does no work for that block. Kernel launch overhead is fixed either - way, so skipping via an ``alive`` flag saved no time. + ``current_level`` is device-side so the same launch can be replayed + until every connected component has been traversed. """ b, i = wp.tid() if b >= num_blocks: return + if reorder_active[b] == int(0): + return n_b = dims[b] if i >= n_b: return vb = vio[b] mb = mio[b] + cur = current_level[b] if level[vb + i] != cur: return @@ -259,9 +266,143 @@ def bfs_step_kernel( if level[vb + j] == int(-1): old = wp.atomic_cas(level, vb + j, int(-1), next_lvl) if old == int(-1): + wp.atomic_max(discovered, b, int(1)) slot = wp.atomic_add(head, b, int(1)) order_buf[vb + slot] = j + @wp.kernel(module=module) + def advance_or_seed_kernel( + num_blocks: int, + dims: wp.array[wp.int32], + vio: wp.array[wp.int32], + degree: wp.array[wp.int32], + level: wp.array[wp.int32], + order_buf: wp.array[wp.int32], + head: wp.array[wp.int32], + current_level: wp.array[wp.int32], + discovered: wp.array[wp.int32], + reorder_active: wp.array[wp.int32], + ): + """Advance a BFS level or seed the next disconnected component.""" + b = wp.tid() + if b >= num_blocks or reorder_active[b] == int(0): + return + n_b = dims[b] + if head[b] >= n_b: + return + + next_level = current_level[b] + int(1) + current_level[b] = next_level + if discovered[b] != int(0): + discovered[b] = int(0) + return + + vb = vio[b] + best_deg = int(2147483647) + best_idx = int(-1) + for i in range(n_b): + if level[vb + i] == int(-1): + d = degree[vb + i] + if d < best_deg: + best_deg = d + best_idx = i + + if best_idx >= int(0): + level[vb + best_idx] = next_level + slot = head[b] + head[b] = slot + int(1) + order_buf[vb + slot] = best_idx + + @wp.kernel(module=module) + def complete_cuda_kernel( + num_blocks: int, + tol: dtype, # type: ignore[valid-type] + A: wp.array[dtype], # type: ignore[valid-type] + dims: wp.array[wp.int32], + mio: wp.array[wp.int32], + vio: wp.array[wp.int32], + degree: wp.array[wp.int32], + level: wp.array[wp.int32], + order_buf: wp.array[wp.int32], + head: wp.array[wp.int32], + reorder_active: wp.array[wp.int32], + ): + """Traverse each block completely inside one persistent CUDA block.""" + tid = wp.tid() + block_dim = wp.block_dim() + lane = tid % block_dim + b = tid / block_dim + if b >= num_blocks or reorder_active[b] == int(0): + return + + n_b = dims[b] + vb = vio[b] + mb = mio[b] + + if lane == int(0): + best_deg = int(2147483647) + best_idx = int(0) + for i in range(n_b): + d = degree[vb + i] + if d < best_deg: + best_deg = d + best_idx = i + level[vb + best_idx] = int(0) + order_buf[vb] = best_idx + head[b] = int(1) + _sync_threads() + + frontier_begin = int(0) + current_level = int(0) + while frontier_begin < n_b: + frontier_end = head[b] + next_level = current_level + int(1) + pos = frontier_begin + while pos < frontier_end: + source = order_buf[vb + pos] + base = mb + source * n_b + j = lane + while j < n_b: + if j != source and wp.abs(A[base + j]) > tol: + wp.atomic_cas(level, vb + j, int(-1), next_level) + j += block_dim + pos += int(1) + _sync_threads() + + chunk_start = int(0) + while chunk_start < n_b: + j = chunk_start + lane + is_new = int(0) + if j < n_b and level[vb + j] == next_level: + is_new = int(1) + prefix = wp.tile_scan_inclusive(wp.tile(is_new)) + write_base = int(0) + if lane == block_dim - int(1): + write_base = wp.atomic_add(head, b, prefix[block_dim - int(1)]) + write_base_tile = wp.tile(write_base) + write_base = write_base_tile[block_dim - int(1)] + if is_new != int(0): + order_buf[vb + write_base + prefix[lane] - int(1)] = j + chunk_start += block_dim + _sync_threads() + + if lane == int(0) and head[b] == frontier_end and frontier_end < n_b: + best_deg = int(2147483647) + best_idx = int(-1) + for i in range(n_b): + if level[vb + i] == int(-1): + d = degree[vb + i] + if d < best_deg: + best_deg = d + best_idx = i + if best_idx >= int(0): + level[vb + best_idx] = next_level + order_buf[vb + frontier_end] = best_idx + head[b] = frontier_end + int(1) + _sync_threads() + frontier_begin = frontier_end + current_level = next_level + @wp.kernel(module=module) def append_unreached_kernel( num_blocks: int, @@ -270,14 +411,14 @@ def append_unreached_kernel( level: wp.array[wp.int32], # type: ignore[valid-type] order_buf: wp.array[wp.int32], # type: ignore[valid-type] head: wp.array[wp.int32], # type: ignore[valid-type] + reorder_active: wp.array[wp.int32], ): - """Launch dims: ``(num_blocks,)``. Appends any vertex with - ``level == -1`` to each block's ``order_buf`` segment in ascending - index order. Serialized per block. - """ + """Append vertices left by an explicitly truncated traversal.""" b = wp.tid() if b >= num_blocks: return + if reorder_active[b] == int(0): + return n_b = dims[b] vb = vio[b] pos = head[b] @@ -294,144 +435,45 @@ def reverse_into_perm_kernel( vio: wp.array[wp.int32], # type: ignore[valid-type] order_buf: wp.array[wp.int32], # type: ignore[valid-type] perm: wp.array[wp.int32], # type: ignore[valid-type] + reorder_active: wp.array[wp.int32], + permutation_valid: wp.array[wp.int32], + permutation_dim: wp.array[wp.int32], ): """Launch dims: ``(num_blocks, max_dim)``. ``perm[i] = order_buf[n-1-i]``.""" b, i = wp.tid() if b >= num_blocks: return + if reorder_active[b] == int(0): + return n_b = dims[b] if i >= n_b: return vb = vio[b] perm[vb + i] = order_buf[vb + (n_b - int(1) - i)] + if i == 0: + permutation_valid[b] = int(1) + permutation_dim[b] = n_b return { + "prepare_reorder": prepare_reorder_kernel, "init_and_degree": init_and_degree_kernel, "select_and_seed": select_and_seed_kernel, "bfs_step": bfs_step_kernel, + "advance_or_seed": advance_or_seed_kernel, + "complete_cuda": complete_cuda_kernel, "append_unreached": append_unreached_kernel, "reverse_into_perm": reverse_into_perm_kernel, } -def _fused_rcm_block_dim(max_dim: int) -> int: - """Pick one CUDA block large enough to assign one thread per vertex.""" - return min(1024, max(32, 1 << (max_dim - 1).bit_length())) - - -@cache -def _make_rcm_batch_fused_tile_kernel(dtype, max_dim: int): - """Create a native-free tiled RCM kernel using shared tiles.""" - module_name = f"rcm_batch_fused_tile_kernels_{getattr(dtype, '__name__', str(dtype))}_{max_dim}" - module = wp.get_module(module_name) - module.options.update({"enable_backward": False, "default_grid_stride": False}) - - @wp.kernel(module=module) - def fused_rcm_tile_kernel( - num_blocks: int, - max_bfs_iters: int, - tol: dtype, # type: ignore[valid-type] - A: wp.array[dtype], # type: ignore[valid-type] - dims: wp.array[wp.int32], - mio: wp.array[wp.int32], - vio: wp.array[wp.int32], - perm: wp.array[wp.int32], - ): - b, lane = wp.tid() - if b >= num_blocks: - return - - n_b = dims[b] - if n_b > max_dim: - return - - mb = mio[b] - vb = vio[b] - - degree = wp.tile_zeros(shape=max_dim, dtype=wp.int32, storage="shared") - level = wp.tile_zeros(shape=max_dim, dtype=wp.int32, storage="shared") - - d = int(0) - if lane < n_b: - row = mb + lane * n_b - for j in range(n_b): - if j == lane: - continue - av = wp.abs(A[row + j]) - if av > tol: - d += int(1) - - wp.tile_scatter_masked(level, lane, int(-1), lane < n_b) - wp.tile_scatter_masked(degree, lane, d, lane < n_b) - - best_idx = int(0) - if lane == 0: - best_deg = int(2147483647) - for i in range(n_b): - deg_i = wp.tile_extract(degree, i) - if deg_i < best_deg: - best_deg = deg_i - best_idx = i - - wp.tile_scatter_masked(level, best_idx, int(0), lane == 0) - - for cur in range(max_bfs_iters): - discovered = wp.bool(False) - # Vertex-owned update: lane ``j`` is the only writer of - # ``level[j]``, so tile_scatter_masked is race-free without CAS. - if lane < n_b and wp.tile_extract(level, lane) == int(-1): - for i in range(n_b): - if wp.tile_extract(level, i) == cur: - av = wp.abs(A[mb + i * n_b + lane]) - if i != lane and av > tol: - discovered = True - wp.tile_scatter_masked(level, lane, cur + int(1), discovered) - - if lane < n_b: - lane_level = wp.tile_extract(level, lane) - cm_pos = int(0) - if lane_level == int(-1): - for i in range(n_b): - if wp.tile_extract(level, i) != int(-1): - cm_pos += int(1) - for i in range(lane): - if wp.tile_extract(level, i) == int(-1): - cm_pos += int(1) - else: - for i in range(n_b): - level_i = wp.tile_extract(level, i) - if level_i != int(-1): - if level_i < lane_level: - cm_pos += int(1) - elif level_i == lane_level and i < lane: - cm_pos += int(1) - - perm[vb + (n_b - int(1) - cm_pos)] = lane - - return fused_rcm_tile_kernel - - # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- -def _default_bfs_iters(max_dim: int) -> int: - """Upper bound on BFS depth we actually execute. - - We use ``max_dim`` (the largest block) as the conservative sizing; all - blocks run the same number of steps in the batched layout. Smaller blocks - just saturate earlier and the subsequent steps become no-ops for them. - - The classical analytic upper bound is ``2*sqrt(n)`` (Cuthill-McKee on a - banded matrix with bandwidth ``sqrt(n)``). We keep that full bound as the - default because lowering it regresses tile-fill on banded-scrambled - matrices at ~5% density (bandwidth ~ 6 but after scrambling BFS needs - several more expansion rounds to recover). Any gains from dropping - launches are then eaten by the extra tiles the factorize/solve kernels - no longer skip. - """ - return 2 * int(math.ceil(math.sqrt(max_dim))) + 4 +def _persistent_block_dim(max_dim: int) -> int: + """Choose enough CUDA lanes to scan matrix rows cooperatively.""" + return min(1024, max(32, 1 << (max_dim - 1).bit_length())) def create_rcm_batch_launch( @@ -446,6 +488,7 @@ def create_rcm_batch_launch( tol: float = 0.0, max_bfs_iters: int | None = None, use_cuda_graph: bool = True, + reuse_permutation: bool = False, device=None, stream=None, ) -> Callable[[], None]: @@ -466,8 +509,10 @@ def create_rcm_batch_launch( memory on replay. num_blocks, max_dim: Host-side sizing used to pick fixed launch dimensions. - tol, max_bfs_iters, use_cuda_graph, device, stream: - Same semantics as :func:`rcm.create_rcm_launch`. + max_bfs_iters: + Optional approximate traversal cap. By default, all components are traversed. + tol, use_cuda_graph, reuse_permutation, device, stream: + Reordering and launch options. """ if perm_flat.dtype != wp.int32: raise TypeError(f"perm_flat must be wp.int32; got {perm_flat.dtype}") @@ -476,41 +521,26 @@ def create_rcm_batch_launch( if device is None: device = A_flat.device device = wp.get_device(device) - if max_bfs_iters is None: - max_bfs_iters = _default_bfs_iters(max_dim) - max_bfs_iters = min(max_bfs_iters, max_dim) - - if dtype == wp.float32 and device.is_cuda and max_dim <= 1024: - fused_kernel = _make_rcm_batch_fused_tile_kernel(dtype, max_dim) - fused_launch = wp.launch_tiled( - fused_kernel, - dim=num_blocks, - inputs=[ - num_blocks, - int(max_bfs_iters), - float(tol), - A_flat, - dims, - mio, - vio, - perm_flat, - ], - device=device, - stream=stream, - block_dim=_fused_rcm_block_dim(max_dim), - record_cmd=True, - ) - - def callback(): - fused_launch.launch() - - if use_cuda_graph: - return create_cuda_graph_callback(callback, device=device, stream=stream) - return callback + complete_traversal = max_bfs_iters is None + if max_bfs_iters is not None: + max_bfs_iters = min(max(0, max_bfs_iters), max_dim) K = _make_rcm_batch_kernels(dtype) - # Pre-record launches with fixed (num_blocks, max_dim) topology. + prepare_reorder_launch = wp.launch( + K["prepare_reorder"], + dim=num_blocks, + inputs=[ + bool(reuse_permutation), + dims, + scratch["permutation_valid"], + scratch["permutation_dim"], + scratch["reorder_active"], + ], + device=device, + stream=stream, + record_cmd=True, + ) init_and_degree_launch = wp.launch( K["init_and_degree"], dim=(num_blocks, max_dim), @@ -524,7 +554,9 @@ def callback(): scratch["degree"], scratch["level"], scratch["head"], - scratch["root"], + scratch["current_level"], + scratch["discovered"], + scratch["reorder_active"], ], device=device, stream=stream, @@ -532,7 +564,45 @@ def callback(): ) select_and_seed_launch = wp.launch( K["select_and_seed"], - dim=(num_blocks,), + dim=num_blocks, + inputs=[ + num_blocks, + dims, + vio, + scratch["degree"], + scratch["level"], + scratch["order_buf"], + scratch["head"], + scratch["reorder_active"], + ], + device=device, + stream=stream, + record_cmd=True, + ) + bfs_step_launch = wp.launch( + K["bfs_step"], + dim=(num_blocks, max_dim), + inputs=[ + num_blocks, + float(tol), + A_flat, + dims, + mio, + vio, + scratch["level"], + scratch["order_buf"], + scratch["head"], + scratch["current_level"], + scratch["discovered"], + scratch["reorder_active"], + ], + device=device, + stream=stream, + record_cmd=True, + ) + advance_or_seed_launch = wp.launch( + K["advance_or_seed"], + dim=num_blocks, inputs=[ num_blocks, dims, @@ -541,43 +611,50 @@ def callback(): scratch["level"], scratch["order_buf"], scratch["head"], - scratch["root"], + scratch["current_level"], + scratch["discovered"], + scratch["reorder_active"], ], device=device, stream=stream, record_cmd=True, ) - # Pre-record one bfs_step launch per iteration, each with its own - # ``cur`` scalar baked in at record time. This removes the need for a - # device-side iteration counter (and its per-step 1-thread increment - # launch) and is race-free by construction: every thread in a given - # launch sees the same compile-time-constant-looking level. - bfs_step_launches = [ - wp.launch( - K["bfs_step"], - dim=(num_blocks, max_dim), + complete_cuda_launch = None + if complete_traversal and device.is_cuda: + block_dim = _persistent_block_dim(max_dim) + complete_cuda_launch = wp.launch( + K["complete_cuda"], + dim=num_blocks * block_dim, inputs=[ num_blocks, - int(cur), float(tol), A_flat, dims, mio, vio, + scratch["degree"], scratch["level"], scratch["order_buf"], scratch["head"], + scratch["reorder_active"], ], device=device, stream=stream, + block_dim=block_dim, record_cmd=True, ) - for cur in range(max_bfs_iters) - ] append_unreached_launch = wp.launch( K["append_unreached"], - dim=(num_blocks,), - inputs=[num_blocks, dims, vio, scratch["level"], scratch["order_buf"], scratch["head"]], + dim=num_blocks, + inputs=[ + num_blocks, + dims, + vio, + scratch["level"], + scratch["order_buf"], + scratch["head"], + scratch["reorder_active"], + ], device=device, stream=stream, record_cmd=True, @@ -585,17 +662,35 @@ def callback(): reverse_launch = wp.launch( K["reverse_into_perm"], dim=(num_blocks, max_dim), - inputs=[num_blocks, dims, vio, scratch["order_buf"], perm_flat], + inputs=[ + num_blocks, + dims, + vio, + scratch["order_buf"], + perm_flat, + scratch["reorder_active"], + scratch["permutation_valid"], + scratch["permutation_dim"], + ], device=device, stream=stream, record_cmd=True, ) + def traverse_level(): + bfs_step_launch.launch() + advance_or_seed_launch.launch() + def callback(): + prepare_reorder_launch.launch() init_and_degree_launch.launch() - select_and_seed_launch.launch() - for step in bfs_step_launches: - step.launch() + if complete_cuda_launch is not None: + complete_cuda_launch.launch() + else: + select_and_seed_launch.launch() + iteration_count = max_dim if complete_traversal else int(max_bfs_iters) + for _ in range(iteration_count): + traverse_level() append_unreached_launch.launch() reverse_launch.launch() diff --git a/newton/_src/solvers/kamino/_src/linalg/linear.py b/newton/_src/solvers/kamino/_src/linalg/linear.py index c0275b5b80..a89999cc4c 100644 --- a/newton/_src/solvers/kamino/_src/linalg/linear.py +++ b/newton/_src/solvers/kamino/_src/linalg/linear.py @@ -571,9 +571,10 @@ class LLTBlockedSolver(DirectSolver[ScalarType, IndexType]): def __init__( self, operator: DenseLinearOperatorData[ScalarType, IndexType] | None = None, - block_size: int = 32, + factorize_block_size: int = 32, + solve_block_size: int = 32, solve_block_dim: int = 128, - factortize_block_dim: int = 128, + factorize_block_dim: int = 128, atol: float | None = None, rtol: float | None = None, ftol: float | None = None, @@ -581,21 +582,47 @@ def __init__( device: wp.DeviceLike | None = None, **kwargs: dict[str, Any], ): + """Initialize a blocked Cholesky solver. + + The factorization and triangular-solve kernels operate on the same + dense ``L`` matrix, but tile their computations independently. Larger + factorization tiles reduce the number of sequential panel steps, while + smaller solve tiles are generally preferable for a single right-hand + side. + + Args: + operator: Linear operator to allocate for immediately, or ``None`` + to defer allocation. + factorize_block_size: Matrix tile width and height used by the + factorization kernel. + solve_block_size: Matrix tile width and height used by the + triangular-solve kernels. + solve_block_dim: Number of threads per CUDA block used to launch + the triangular-solve kernels. + factorize_block_dim: Number of threads per CUDA block used to + launch the factorization kernel. + atol: Absolute solve tolerance. + rtol: Relative solve tolerance. + ftol: Factorization tolerance. + dtype: Scalar data type. + device: Device on which to allocate and run the solver. + **kwargs: Additional arguments forwarded to :class:`DirectSolver`. + """ # Declare LLT-specific internal data self._L: wp.array[ScalarType] | None = None """A flat array containing the Cholesky factorization of each matrix block.""" self._y: wp.array[ScalarType] | None = None """A flat array containing the intermediate results for the solve operation.""" - # Cache the fixed block size - self._block_size: int = block_size + self._factorize_block_size: int = factorize_block_size + self._solve_block_size: int = solve_block_size self._solve_block_dim: int = solve_block_dim - self._factortize_block_dim: int = factortize_block_dim + self._factorize_block_dim: int = factorize_block_dim # Create the factorization and solve kernels - self._factorize_kernel = factorize.make_llt_blocked_factorize_kernel(block_size) - self._solve_kernel = factorize.make_llt_blocked_solve_kernel(block_size) - self._solve_inplace_kernel = factorize.make_llt_blocked_solve_inplace_kernel(block_size) + self._factorize_kernel = factorize.make_llt_blocked_factorize_kernel(self._factorize_block_size) + self._solve_kernel = factorize.make_llt_blocked_solve_kernel(self._solve_block_size) + self._solve_inplace_kernel = factorize.make_llt_blocked_solve_inplace_kernel(self._solve_block_size) # Initialize base class members super().__init__( @@ -655,7 +682,7 @@ def _factorize_impl(self, A: wp.array[ScalarType]) -> None: factorize.llt_blocked_factorize( kernel=self._factorize_kernel, num_blocks=self._operator.info.num_blocks, - block_dim=self._factortize_block_dim, + block_dim=self._factorize_block_dim, dim=self._operator.info.dim, mio=self._operator.info.mio, A=A, @@ -749,7 +776,7 @@ def _allocate_impl(self, operator, **kwargs: dict[str, Any]) -> None: maxiter=self._maxiter, Mi=self._Mi, callback=None, - use_cuda_graph=True, + use_graph=True, use_graph_conditionals=self._use_graph_conditionals, loop_granularity=self.loop_granularity, ) @@ -763,7 +790,7 @@ def _allocate_impl(self, operator, **kwargs: dict[str, Any]) -> None: maxiter=self._maxiter, Mi=self._Mi, callback=None, - use_cuda_graph=True, + use_graph=True, use_graph_conditionals=self._use_graph_conditionals, loop_granularity=self.loop_granularity, ) @@ -861,7 +888,7 @@ def _allocate_impl(self, operator, **kwargs: dict[str, Any]) -> None: maxiter=self._maxiter, Mi=self._Mi, callback=None, - use_cuda_graph=True, + use_graph=True, use_graph_conditionals=self._use_graph_conditionals, loop_granularity=self.loop_granularity, ) @@ -875,7 +902,7 @@ def _allocate_impl(self, operator, **kwargs: dict[str, Any]) -> None: maxiter=self._maxiter, Mi=self._Mi, callback=None, - use_cuda_graph=True, + use_graph=True, use_graph_conditionals=self._use_graph_conditionals, loop_granularity=self.loop_granularity, ) diff --git a/newton/_src/solvers/kamino/_src/models/builders/basics.py b/newton/_src/solvers/kamino/_src/models/builders/basics.py index 51974702d1..09d71935a8 100644 --- a/newton/_src/solvers/kamino/_src/models/builders/basics.py +++ b/newton/_src/solvers/kamino/_src/models/builders/basics.py @@ -20,7 +20,7 @@ from ......core.types import Axis from ...core import ModelBuilderKamino, inertia -from ...core.joints import JointActuationType, JointDoFType +from ...core.joints import JOINT_TAUMAX, JointActuationType, JointDoFType from ...core.math import FLOAT32_MAX, FLOAT32_MIN, I_3, axis_to_mat33 from ...core.shapes import BoxShape, PlaneShape, SphereShape @@ -174,6 +174,7 @@ def build_box_pendulum( B_r_Bj=wp.vec3f(0.0, 0.0, 0.5 * h + z_0), F_r_Fj=wp.vec3f(-0.5 * d, 0.0, 0.0), X_Bj=axis_to_mat33(Axis.Y), + tau_j_max=math.inf, # Setting effort limit to match USD convention (`inf` for active joints) a_j=1.0 if dynamic_joints else None, b_j=0.1 if dynamic_joints else None, k_p_j=100.0 if implicit_pd else None, @@ -513,6 +514,7 @@ def build_boxes_hinged( B_r_Bj=wp.vec3f(0.25, 0.05, 0.0), F_r_Fj=wp.vec3f(-0.25, -0.05, 0.0), X_Bj=axis_to_mat33(Axis.Y), + tau_j_max=math.inf, # Setting effort limit to match USD convention (`inf` for active joints) a_j=1.0 if dynamic_joints else None, b_j=0.1 if dynamic_joints else None, k_p_j=100.0 if implicit_pd else None, @@ -1085,6 +1087,9 @@ def build_boxes_fourbar( X_Bj=X_j, q_j_min=[qmin], q_j_max=[qmax], + # Setting effort limit to match USD convention (`inf` sentinel for + # active joints, `JOINT_TAUMAX` for passive joints). + tau_j_max=math.inf if 1 in actuator_ids else JOINT_TAUMAX, a_j=0.1 if dynamic_joints else None, b_j=0.001 if dynamic_joints else None, k_p_j=1000.0 if implicit_pd else None, @@ -1103,6 +1108,9 @@ def build_boxes_fourbar( X_Bj=X_j, q_j_min=[qmin], q_j_max=[qmax], + # Setting effort limit to match USD convention (`inf` sentinel for + # active joints, `JOINT_TAUMAX` for passive joints). + tau_j_max=math.inf if 2 in actuator_ids else JOINT_TAUMAX, world_index=world_index, ) @@ -1117,6 +1125,9 @@ def build_boxes_fourbar( X_Bj=X_j, q_j_min=[qmin], q_j_max=[qmax], + # Setting effort limit to match USD convention (`inf` sentinel for + # active joints, `JOINT_TAUMAX` for passive joints). + tau_j_max=math.inf if 3 in actuator_ids else JOINT_TAUMAX, world_index=world_index, ) @@ -1131,6 +1142,9 @@ def build_boxes_fourbar( X_Bj=X_j, q_j_min=[qmin], q_j_max=[qmax], + # Setting effort limit to match USD convention (`inf` sentinel for + # active joints, `JOINT_TAUMAX` for passive joints). + tau_j_max=math.inf if 4 in actuator_ids else JOINT_TAUMAX, world_index=world_index, ) diff --git a/newton/_src/solvers/kamino/_src/models/builders/testing.py b/newton/_src/solvers/kamino/_src/models/builders/testing.py index 5da9fcff4b..24f677e807 100644 --- a/newton/_src/solvers/kamino/_src/models/builders/testing.py +++ b/newton/_src/solvers/kamino/_src/models/builders/testing.py @@ -19,7 +19,7 @@ from ......core.types import Axis from ...core import ModelBuilderKamino from ...core.joints import JointActuationType, JointDoFType -from ...core.math import I_3, axis_to_mat33, quat_from_euler_xyz +from ...core.math import I_3, axis_to_mat33 from ...core.shapes import ( BoxShape, CapsuleShape, @@ -1363,6 +1363,7 @@ def alter_binary_joint( base_joint = copy.deepcopy(builder.joints[0][0]) if make_floating_base: base_joint.dof_type = JointDoFType.FREE + base_joint.__post_init__() # Will correctly populate joint dynamics etc. builder_alt.add_joint_descriptor(base_joint) joint = copy.deepcopy(builder.joints[0][1]) if make_actuated: @@ -1392,10 +1393,7 @@ def make_unary(builder: ModelBuilderKamino) -> ModelBuilderKamino: geom_.shape = builder.shapes[geom.uid] geom_.body = geom.body - 1 if geom_.body == -1: - # wp.transform_set_translation(geom_.offset, body_0_offset) - geom_.offset[0] = body_0_offset[0] - geom_.offset[1] = body_0_offset[1] - geom_.offset[2] = body_0_offset[2] + wp.transform_set_translation(geom_.offset, body_0_offset) builder_unary.add_geometry_descriptor(geom_) return builder_unary @@ -1631,11 +1629,11 @@ def make_single_shape_pair_builder( # Compute bottom box position and orientation r_b = wp.vec3f(bottom_xyz) - r_dz - q_b = quat_from_euler_xyz(wp.vec3f(*bottom_rpy)) + q_b = wp.quat_from_euler(wp.vec3f(*bottom_rpy), 0, 1, 2) # Compute top sphere position and orientation r_t = wp.vec3f(top_xyz) + r_dz - q_t = quat_from_euler_xyz(wp.vec3f(*top_rpy)) + q_t = wp.quat_from_euler(wp.vec3f(*top_rpy), 0, 1, 2) # Create the shape descriptors for bottom and top shapes # with special handling for PlaneShape diff --git a/newton/_src/solvers/kamino/_src/solver_kamino_impl.py b/newton/_src/solvers/kamino/_src/solver_kamino_impl.py index 405bbb4866..6795ebccfb 100644 --- a/newton/_src/solvers/kamino/_src/solver_kamino_impl.py +++ b/newton/_src/solvers/kamino/_src/solver_kamino_impl.py @@ -57,9 +57,11 @@ set_floating_base, ) from .linalg import ConjugateResidualSolver, IterativeSolver, LinearSolverNameToType +from .solvers.common import WarmStartMode +from .solvers.dvi import DVISolver from .solvers.fk import ForwardKinematicsSolver from .solvers.metrics import SolutionMetrics -from .solvers.padmm import PADMMSolver, PADMMWarmStartMode +from .solvers.padmm import PADMMSolver from .solvers.warmstart import WarmstarterContacts, WarmstarterLimits from .utils import logger as msg @@ -91,10 +93,9 @@ class SolverKaminoImpl(SolverBase): Config = SolverKamino.Config """ - Defines a type alias of the PADMM solver configurations container, including convergence - criteria, maximum iterations, and options for the linear solver and preconditioning. - - See :class:`PADMMSolverConfig` for the full list of configuration options and their descriptions. + Defines a type alias of the public Kamino solver configuration container, + including the selected forward-dynamics solver, convergence criteria, and + options for the linear solver and preconditioning. """ ResetCallbackType = Callable[["SolverKaminoImpl", StateKamino], None] @@ -151,7 +152,16 @@ def __init__( # Cache the solver config and parse relevant options for internal use self._config: SolverKaminoImpl.Config = config - self._warmstart_mode: PADMMWarmStartMode = PADMMWarmStartMode.from_string(config.padmm.warmstart_mode) + if config.dynamics_solver == "padmm": + warmstart_mode = config.padmm.warmstart_mode + contact_warmstart_method = config.padmm.contact_warmstart_method + elif config.dynamics_solver == "dvi": + warmstart_mode = config.dvi.warmstart_mode + contact_warmstart_method = config.dvi.contact_warmstart_method + else: + raise ValueError(f"Unsupported dynamics solver: {config.dynamics_solver}") + self._warmstart_mode = WarmStartMode.from_string(warmstart_mode) + self._contact_warmstart_method = WarmstarterContacts.Method.from_string(contact_warmstart_method) self._rotation_correction: JointCorrectionMode = JointCorrectionMode.from_string(config.rotation_correction) # --------------------------------------------------------------------------- @@ -229,14 +239,31 @@ def __init__( ) # Allocate the forward dynamics solver on the device - self._solver_fd = PADMMSolver( - model=self._model, - config=self._config.padmm, - warmstart=self._warmstart_mode, - use_acceleration=self._config.padmm.use_acceleration, - use_graph_conditionals=self._config.padmm.use_graph_conditionals, - collect_info=self._config.collect_solver_info, - ) + if self._config.dynamics_solver == "padmm": + self._solver_fd = PADMMSolver( + model=self._model, + config=self._config.padmm, + warmstart=self._warmstart_mode, + use_acceleration=self._config.padmm.use_acceleration, + use_graph_conditionals=self._config.padmm.use_graph_conditionals, + collect_info=self._config.collect_solver_info, + ) + elif self._config.dynamics_solver == "dvi": + # DVI consumes Kamino's unified joint, limit, and contact + # DualProblem rather than rebuilding standalone constraint pipelines. + self._solver_fd = DVISolver( + model=self._model, + data=self._data, + limits=self._limits, + contacts=contacts, + jacobians=self._jacobians if isinstance(self._jacobians, SparseSystemJacobians) else None, + problem=self._problem_fd, + config=self._config.dvi, + warmstart=self._warmstart_mode, + collect_info=self._config.collect_solver_info, + ) + else: + raise ValueError(f"Unsupported dynamics solver: {self._config.dynamics_solver}") # Allocate the forward kinematics solver on the device self._solver_fk = None @@ -255,7 +282,6 @@ def __init__( # Allocate additional internal data for reset operations with wp.ScopedDevice(self._model.device): - self._all_worlds_mask = wp.ones(shape=(self._model.size.num_worlds,), dtype=wp.bool) self._base_q = wp.zeros(shape=(self._model.size.num_worlds,), dtype=wp.transformf) self._base_u = wp.zeros(shape=(self._model.size.num_worlds,), dtype=wp.spatial_vectorf) self._bodies_u_zeros = wp.zeros(shape=(self._model.size.sum_of_num_bodies,), dtype=wp.spatial_vectorf) @@ -265,11 +291,11 @@ def __init__( # Allocate the contacts warmstarter if enabled self._ws_limits: WarmstarterLimits | None = None self._ws_contacts: WarmstarterContacts | None = None - if self._warmstart_mode == PADMMWarmStartMode.CONTAINERS: + if self._warmstart_mode == WarmStartMode.CONTAINERS: self._ws_limits = WarmstarterLimits(limits=self._limits) self._ws_contacts = WarmstarterContacts( contacts=contacts, - method=WarmstarterContacts.Method.from_string(self._config.padmm.contact_warmstart_method), + method=self._contact_warmstart_method, ) # Allocate the solution metrics evaluator if enabled @@ -321,7 +347,7 @@ def problem_fd(self) -> DualProblem: return self._problem_fd @property - def solver_fd(self) -> PADMMSolver: + def solver_fd(self) -> PADMMSolver | DVISolver: """ Returns the forward dynamics solver. """ @@ -414,8 +440,7 @@ def _check_length(data: wp.array[Any], name: str, expected: int): if data is not None and data.shape[0] != expected: raise ValueError(f"Invalid shape for {name}: Expected ({expected},), but got {data.shape}.") - # Resolve and validate world mask - world_mask = self._all_worlds_mask if world_mask is None else world_mask + # Validate world mask _check_length(world_mask, "world_mask", self._model.size.num_worlds) # Resolve and validate reset config @@ -493,11 +518,11 @@ def _check_length(data: wp.array[Any], name: str, expected: int): # Extract joint state into pre-allocated actuator state buffers extract_actuators_state_from_joints( model=self._model, - world_mask=world_mask, joint_q=joint_q if joint_q is not None else state.q_j, joint_u=joint_u if joint_u is not None else state.dq_j, actuator_q=self._actuators_q, actuator_u=self._actuators_u, + world_mask=world_mask, ) actuator_q = self._actuators_q if joint_q is not None else actuator_q actuator_u = self._actuators_u if joint_u is not None else actuator_u @@ -631,8 +656,10 @@ def _check_length(data: wp.array[Any], name: str, expected: int): # Currently, only the position-level (iterative) FK solve can fail if actuator_q is not None: wp.copy(success_mask, self._solver_fk.newton_success) - else: + elif world_mask is not None: wp.copy(success_mask, world_mask) + else: + success_mask.fill_(True) # Run the post-reset callback if it has been set self._run_post_reset_callback(state_out=state) @@ -702,7 +729,13 @@ def step( @override def notify_model_changed(self, flags: ModelFlags | int) -> None: - pass # TODO: Migrate implementation when we fully integrate with Newton + if self._solver_fk is not None: + self._solver_fk.notify_model_changed(flags) + + def validate_model_changed(self, flags: ModelFlags | int) -> None: + """Validate solver-specific structural invariants before model updates.""" + if self._solver_fk is not None: + self._solver_fk.validate_model_changed(flags) @override def update_contacts(self, contacts: Contacts, state: State | None = None) -> None: @@ -867,7 +900,7 @@ def _reset_solver_data(self, world_mask: wp.array[wp.bool] | None = None): self._solver_fd.reset(problem=self._problem_fd, world_mask=world_mask) # Reset the warm-starting caches if enabled - if self._warmstart_mode == PADMMWarmStartMode.CONTAINERS: + if self._warmstart_mode == WarmStartMode.CONTAINERS: self._ws_limits.reset(world_mask=world_mask) self._ws_contacts.reset(world_mask=world_mask) @@ -956,8 +989,8 @@ def _update_constraints(self, contacts: ContactsKamino | None = None): """ # If warm-starting is enabled, initialize unilateral # constraints containers from the current solver data - if self._warmstart_mode > PADMMWarmStartMode.NONE: - if self._warmstart_mode == PADMMWarmStartMode.CONTAINERS: + if self._warmstart_mode > WarmStartMode.NONE: + if self._warmstart_mode == WarmStartMode.CONTAINERS: self._ws_limits.warmstart(self._limits) self._ws_contacts.warmstart(self._model, self._data, contacts) self._solver_fd.warmstart( @@ -1000,7 +1033,7 @@ def _update_constraints(self, contacts: ContactsKamino | None = None): # If warmstarting is enabled, update the limits and contacts caches # with the constraint reactions generated by the dynamics solver # NOTE: This needs to happen after unpacking the multipliers - if self._warmstart_mode == PADMMWarmStartMode.CONTAINERS: + if self._warmstart_mode == WarmStartMode.CONTAINERS: self._ws_limits.update(self._limits) self._ws_contacts.update(contacts) diff --git a/newton/_src/solvers/kamino/_src/solvers/__init__.py b/newton/_src/solvers/kamino/_src/solvers/__init__.py index 788a9b3e2b..24e7e6afff 100644 --- a/newton/_src/solvers/kamino/_src/solvers/__init__.py +++ b/newton/_src/solvers/kamino/_src/solvers/__init__.py @@ -3,6 +3,7 @@ """Numerical Solvers for Constraint Rigid Multi-Body Kinematics & Dynamics""" +from .dvi import DVISolver from .fk import ForwardKinematicsSolver from .padmm import PADMMSolver, PADMMWarmStartMode @@ -11,6 +12,7 @@ ### __all__ = [ + "DVISolver", "ForwardKinematicsSolver", "PADMMSolver", "PADMMWarmStartMode", diff --git a/newton/_src/solvers/kamino/_src/solvers/common.py b/newton/_src/solvers/kamino/_src/solvers/common.py new file mode 100644 index 0000000000..8092b6999d --- /dev/null +++ b/newton/_src/solvers/kamino/_src/solvers/common.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Shared data and kernels for Kamino dual forward-dynamics solvers.""" + +from __future__ import annotations + +from enum import IntEnum +from typing import Any + +import warp as wp + +from ..core.size import SizeKamino + + +class WarmStartMode(IntEnum): + """Warm-start source used by a Kamino dual solver.""" + + NONE = -1 + """Start without cached solution data.""" + + INTERNAL = 0 + """Initialize from the solver's internally cached solution.""" + + CONTAINERS = 1 + """Initialize from external joint, limit, and contact containers.""" + + @classmethod + def from_string(cls, value: str) -> WarmStartMode: + """Convert a string to a warm-start mode.""" + try: + return cls[value.upper()] + except KeyError as error: + raise ValueError( + f"Invalid WarmStartMode: {value}. Valid options are: {[mode.name for mode in cls]}" + ) from error + + @staticmethod + def parse_usd_attribute(value: str, context: dict[str, Any] | None = None) -> str: + """Parse a warm-start option imported from USD.""" + del context + if not isinstance(value, str): + raise TypeError("Parser expects input of type 'str'.") + value = value.lower().strip() + if value not in {"none", "internal", "containers"}: + raise ValueError(f"Warmstart parameter '{value}' is not a valid option.") + return value + + +class DualSolution: + """Constraint impulse and post-event velocity arrays for a dual solver.""" + + def __init__(self, size: SizeKamino | None = None): + self.lambdas: wp.array[wp.float32] | None = None + """Constraint impulses, shape ``(sum_of_max_total_cts,)``.""" + self.v_plus: wp.array[wp.float32] | None = None + """Post-event constraint-space velocities, shape ``(sum_of_max_total_cts,)``.""" + if size is not None: + self.finalize(size) + + def finalize(self, size: SizeKamino) -> None: + """Allocate solution arrays for a model size.""" + self.lambdas = wp.zeros(size.sum_of_max_total_cts, dtype=wp.float32) + self.v_plus = wp.zeros(size.sum_of_max_total_cts, dtype=wp.float32) + + def zero(self) -> None: + """Reset all solution arrays to zero.""" + self.lambdas.zero_() + self.v_plus.zero_() + + +@wp.kernel +def warmstart_joint_constraints( + model_time_dt: wp.array[wp.float32], + joint_wid: wp.array[wp.int32], + joint_num_dynamic_cts: wp.array[wp.int32], + joint_num_kinematic_cts: wp.array[wp.int32], + joint_dynamic_cts_offset_joint_cts: wp.array[wp.int32], + joint_kinematic_cts_offset_joint_cts: wp.array[wp.int32], + joint_dynamic_cts_offset_total_cts: wp.array[wp.int32], + joint_kinematic_cts_offset_total_cts: wp.array[wp.int32], + joint_lambda_j: wp.array[wp.float32], + problem_P: wp.array[wp.float32], + x_0: wp.array[wp.float32], + y_0: wp.array[wp.float32], + z_0: wp.array[wp.float32], +): + """Initialize bilateral constraint iterates from cached joint reactions.""" + jid = wp.tid() + wid_j = joint_wid[jid] + num_dynamic_cts_j = joint_num_dynamic_cts[jid] + num_kinematic_cts_j = joint_num_kinematic_cts[jid] + dt = model_time_dt[wid_j] + joint_dyn_cts_start = joint_dynamic_cts_offset_joint_cts[jid] + joint_kin_cts_start = joint_kinematic_cts_offset_joint_cts[jid] + dyn_cts_row_start_j = joint_dynamic_cts_offset_total_cts[jid] + kin_cts_row_start_j = joint_kinematic_cts_offset_total_cts[jid] + + # Convert cached forces to preconditioned impulses. Joint constraints do + # not cache a post-event velocity, so their dual iterate starts at zero. + for j in range(num_dynamic_cts_j): + P_j = problem_P[dyn_cts_row_start_j + j] + lambda_j = (dt / P_j) * joint_lambda_j[joint_dyn_cts_start + j] + x_0[dyn_cts_row_start_j + j] = lambda_j + y_0[dyn_cts_row_start_j + j] = lambda_j + z_0[dyn_cts_row_start_j + j] = 0.0 + for j in range(num_kinematic_cts_j): + P_j = problem_P[kin_cts_row_start_j + j] + lambda_j = (dt / P_j) * joint_lambda_j[joint_kin_cts_start + j] + x_0[kin_cts_row_start_j + j] = lambda_j + y_0[kin_cts_row_start_j + j] = lambda_j + z_0[kin_cts_row_start_j + j] = 0.0 + + +@wp.kernel +def warmstart_limit_constraints( + model_time_dt: wp.array[wp.float32], + model_info_total_cts_offset: wp.array[wp.int32], + data_info_limit_cts_group_offset: wp.array[wp.int32], + limit_model_num_active: wp.array[wp.int32], + limit_wid: wp.array[wp.int32], + limit_lid: wp.array[wp.int32], + limit_reaction: wp.array[wp.float32], + limit_velocity: wp.array[wp.float32], + problem_P: wp.array[wp.float32], + x_0: wp.array[wp.float32], + y_0: wp.array[wp.float32], + z_0: wp.array[wp.float32], +): + """Initialize limit constraint iterates from cached limit data.""" + lid = wp.tid() + if lid >= limit_model_num_active[0]: + return + + wid = limit_wid[lid] + vio_l = model_info_total_cts_offset[wid] + data_info_limit_cts_group_offset[wid] + limit_lid[lid] + P_l = problem_P[vio_l] + # Reactions are cached as forces and velocities in physical units. + lambda_l = limit_reaction[lid] * model_time_dt[wid] / P_l + v_plus_l = limit_velocity[lid] * P_l + x_0[vio_l] = lambda_l + y_0[vio_l] = lambda_l + z_0[vio_l] = v_plus_l + + +@wp.kernel +def warmstart_contact_constraints( + model_time_dt: wp.array[wp.float32], + model_info_total_cts_offset: wp.array[wp.int32], + data_info_contact_cts_group_offset: wp.array[wp.int32], + contact_model_num_contacts: wp.array[wp.int32], + contact_wid: wp.array[wp.int32], + contact_cid: wp.array[wp.int32], + contact_material: wp.array[wp.vec2f], + contact_reaction: wp.array[wp.vec3f], + contact_velocity: wp.array[wp.vec3f], + problem_P: wp.array[wp.float32], + x_0: wp.array[wp.float32], + y_0: wp.array[wp.float32], + z_0: wp.array[wp.float32], +): + """Initialize contact constraint iterates from cached contact data.""" + cid = wp.tid() + if cid >= contact_model_num_contacts[0]: + return + + wid = contact_wid[cid] + vio_k = model_info_total_cts_offset[wid] + data_info_contact_cts_group_offset[wid] + 3 * contact_cid[cid] + P_k = problem_P[vio_k] + lambda_k = contact_reaction[cid] * model_time_dt[wid] / P_k + v_plus_k = contact_velocity[cid] * P_k + mu_k = contact_material[cid][0] + # Apply the De Saxce correction to recover the solver's dual variable. + v_plus_k.z += mu_k * wp.sqrt(v_plus_k.x * v_plus_k.x + v_plus_k.y * v_plus_k.y) + + for k in range(3): + x_0[vio_k + k] = lambda_k[k] + y_0[vio_k + k] = lambda_k[k] + z_0[vio_k + k] = v_plus_k[k] + + +@wp.kernel +def apply_dual_preconditioner_to_solution( + problem_dim: wp.array[wp.int32], + problem_vio: wp.array[wp.int32], + problem_P: wp.array[wp.float32], + solution_lambdas: wp.array[wp.float32], + solution_v_plus: wp.array[wp.float32], +): + """Convert cached physical solution values to preconditioned solver units.""" + wid, tid = wp.tid() + if tid >= problem_dim[wid]: + return + v_i = problem_vio[wid] + tid + P_i = problem_P[v_i] + # Constraint impulses and velocities scale inversely under dual preconditioning. + solution_lambdas[v_i] /= P_i + solution_v_plus[v_i] *= P_i diff --git a/newton/_src/solvers/kamino/_src/solvers/dvi/__init__.py b/newton/_src/solvers/kamino/_src/solvers/dvi/__init__.py new file mode 100644 index 0000000000..ae90fdd8ce --- /dev/null +++ b/newton/_src/solvers/kamino/_src/solvers/dvi/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""DVI solver for Kamino's dual forward-dynamics problem.""" + +from .solver import DVISolver +from .types import DVIConfigStruct, DVIData, DVIInfo, DVIState, DVIStatus, convert_config_to_struct + +__all__ = [ + "DVIConfigStruct", + "DVIData", + "DVIInfo", + "DVISolver", + "DVIState", + "DVIStatus", + "convert_config_to_struct", +] diff --git a/newton/_src/solvers/kamino/_src/solvers/dvi/kernels.py b/newton/_src/solvers/kamino/_src/solvers/dvi/kernels.py new file mode 100644 index 0000000000..7030b9c20b --- /dev/null +++ b/newton/_src/solvers/kamino/_src/solvers/dvi/kernels.py @@ -0,0 +1,1041 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Warp kernels for the Kamino DVI solver.""" + +from __future__ import annotations + +import warp as wp + +from ...core.math import FLOAT32_EPS +from ..padmm.math import project_to_coulomb_cone, project_to_coulomb_dual_cone +from .projections import ( + contact_trace_preconditioner as _contact_trace_preconditioner, +) +from .projections import ( + project_contact_block_update as _project_contact_block_update, +) +from .projections import ( + project_contact_diagonal_update as _project_contact_diagonal_update, +) +from .types import DVIConfigStruct, DVIStatus + +wp.set_module_options({"enable_backward": False}) + +float32 = wp.float32 +int32 = wp.int32 +mat33f = wp.mat33f +vec3f = wp.vec3f + + +@wp.func +def _compute_row_velocity( + ncts: int32, + mio: int32, + vio: int32, + row: int32, + D: wp.array[float32], + v_f: wp.array[float32], + lambdas: wp.array[float32], +) -> float32: + # Full constraint-space velocity of one row: v_f[row] + sum_j D[row, j] * lambda[j]. + # The sum spans all columns, so a unilateral row picks up the D_ub * lambda_b + # contribution from joint impulses (and a joint row picks up D_bu * lambda_u). + v = v_f[vio + row] + m_i = mio + ncts * row + for j in range(ncts): + v += D[m_i + j] * lambdas[vio + j] + return v + + +@wp.func +def _contact_velocity_aug( + ncts: int32, + mio: int32, + vio: int32, + ccgo: int32, + cio: int32, + cid: int32, + D: wp.array[float32], + v_f: wp.array[float32], + lambdas: wp.array[float32], + mu: wp.array[float32], +) -> vec3f: + # Contact rows are [t0, t1, n]. De Saxce augments the normal velocity by + # mu * ||v_t|| before enforcing Coulomb-cone complementarity. + ccio = ccgo + 3 * cid + v_t0 = _compute_row_velocity(ncts, mio, vio, ccio + 0, D, v_f, lambdas) + v_t1 = _compute_row_velocity(ncts, mio, vio, ccio + 1, D, v_f, lambdas) + v_n = _compute_row_velocity(ncts, mio, vio, ccio + 2, D, v_f, lambdas) + vt_norm = wp.sqrt(v_t0 * v_t0 + v_t1 * v_t1) + return vec3f(v_t0, v_t1, v_n + mu[cio + cid] * vt_norm) + + +@wp.kernel +def _reset_dvi_solver_data( + # Inputs: + world_mask: wp.array[wp.bool], + problem_vio: wp.array[int32], + problem_maxdim: wp.array[int32], + # Outputs: + solution_lambdas: wp.array[float32], + solution_v_plus: wp.array[float32], +): + wid, tid = wp.tid() + if not world_mask[wid] or tid >= problem_maxdim[wid]: + return + v_i = problem_vio[wid] + tid + solution_lambdas[v_i] = 0.0 + solution_v_plus[v_i] = 0.0 + + +@wp.kernel +def _reset_dvi_status( + # Outputs: + solver_status: wp.array[DVIStatus], +): + wid = wp.tid() + solver_status[wid] = DVIStatus() + + +@wp.kernel +def _copy_bilateral_block( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_njc: wp.array[int32], + problem_D: wp.array[float32], + bilateral_mio: wp.array[int32], + bilateral_vio: wp.array[int32], + # Outputs: + bilateral_D: wp.array[float32], + bilateral_P: wp.array[float32], +): + wid, tid = wp.tid() + + njc = problem_njc[wid] + if njc == 0 or tid >= njc * njc: + return + + ncts = problem_dim[wid] + pmio = problem_mio[wid] + bmio = bilateral_mio[wid] + bvio = bilateral_vio[wid] + row = tid // njc + col = tid - row * njc + + D_rr = problem_D[pmio + ncts * row + row] + D_cc = problem_D[pmio + ncts * col + col] + p_row = wp.sqrt(1.0 / (wp.abs(D_rr) + FLOAT32_EPS)) + p_col = wp.sqrt(1.0 / (wp.abs(D_cc) + FLOAT32_EPS)) + + val = p_row * problem_D[pmio + ncts * row + col] * p_col + if row == col: + # Smaller floors reduce equality residual, but closed-loop robots lose contact below this. + val += float32(7.0e-7) + bilateral_P[bvio + row] = p_row + bilateral_D[bmio + njc * row + col] = val + + +@wp.kernel +def _compute_dvi_contact_block_inverse( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_D: wp.array[float32], + solver_config: wp.array[DVIConfigStruct], + # Outputs: + contact_block_inv: wp.array[mat33f], +): + wid, cid = wp.tid() + + nc = problem_nc[wid] + if cid >= nc: + return + + ncts = problem_dim[wid] + mio = problem_mio[wid] + ccgo = problem_ccgo[wid] + cio = problem_cio[wid] + ccio = ccgo + int32(3) * cid + cfg = solver_config[wid] + D_inv = mat33f(0.0) + + if not cfg.contact_block_preconditioner: + contact_block_inv[cio + cid] = D_inv + return + + r0 = mio + ncts * (ccio + 0) + r1 = mio + ncts * (ccio + 1) + r2 = mio + ncts * (ccio + 2) + + d00 = problem_D[r0 + ccio + 0] + d01 = float32(0.5) * (problem_D[r0 + ccio + 1] + problem_D[r1 + ccio + 0]) + d02 = float32(0.5) * (problem_D[r0 + ccio + 2] + problem_D[r2 + ccio + 0]) + d11 = problem_D[r1 + ccio + 1] + d12 = float32(0.5) * (problem_D[r1 + ccio + 2] + problem_D[r2 + ccio + 1]) + d22 = problem_D[r2 + ccio + 2] + + diag_max = wp.max(wp.max(wp.abs(d00), wp.abs(d11)), wp.abs(d22)) + if diag_max > FLOAT32_EPS: + D_reg = mat33f( + d00 + cfg.regularization, + d01, + d02, + d01, + d11 + cfg.regularization, + d12, + d02, + d12, + d22 + cfg.regularization, + ) + det = wp.determinant(D_reg) + det_min = FLOAT32_EPS * diag_max * diag_max * diag_max + if det > det_min: + D_inv = wp.inverse(D_reg) + + contact_block_inv[cio + cid] = D_inv + + +@wp.kernel +def _build_bilateral_rhs( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_vio: wp.array[int32], + problem_njc: wp.array[int32], + problem_D: wp.array[float32], + problem_v_f: wp.array[float32], + bilateral_vio: wp.array[int32], + bilateral_P: wp.array[float32], + solution_lambdas: wp.array[float32], + # Outputs: + bilateral_rhs: wp.array[float32], +): + wid, row = wp.tid() + + njc = problem_njc[wid] + if row >= njc: + return + + ncts = problem_dim[wid] + pmio = problem_mio[wid] + pvio = problem_vio[wid] + bvio = bilateral_vio[wid] + + # Columns njc..ncts are the unilateral rows, so this loop subtracts the + # D_bu * lambda_u coupling: the current limit and contact impulses enter the + # joint solve, yielding rhs = -(v_f,b + D_bu * lambda_u). + rhs = -problem_v_f[pvio + row] + for col in range(njc, ncts): + rhs -= problem_D[pmio + ncts * row + col] * solution_lambdas[pvio + col] + bilateral_rhs[bvio + row] = bilateral_P[bvio + row] * rhs + + +@wp.kernel +def _scatter_bilateral_solution( + # Inputs: + problem_vio: wp.array[int32], + problem_njc: wp.array[int32], + bilateral_vio: wp.array[int32], + bilateral_P: wp.array[float32], + bilateral_solution: wp.array[float32], + # Outputs: + solution_lambdas: wp.array[float32], +): + wid, row = wp.tid() + + njc = problem_njc[wid] + if row >= njc: + return + + bvio = bilateral_vio[wid] + solution_lambdas[problem_vio[wid] + row] = bilateral_P[bvio + row] * bilateral_solution[bvio + row] + + +@wp.kernel +def _compute_dvi_status_residuals( + # Inputs: + problem_dim: wp.array[int32], + problem_vio: wp.array[int32], + problem_njc: wp.array[int32], + problem_nl: wp.array[int32], + problem_nc: wp.array[int32], + problem_lcgo: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + solver_config: wp.array[DVIConfigStruct], + state_v_aug: wp.array[float32], + solution_lambdas: wp.array[float32], + # Outputs: + solver_status: wp.array[DVIStatus], +): + wid = wp.tid() + + ncts = problem_dim[wid] + vio = problem_vio[wid] + njc = problem_njc[wid] + nl = problem_nl[wid] + nc = problem_nc[wid] + lcgo = problem_lcgo[wid] + ccgo = problem_ccgo[wid] + cio = problem_cio[wid] + cfg = solver_config[wid] + + status = solver_status[wid] + if status.iterations == 0: + status.iterations = int32(1) + + # These terminal diagnostics are distinct from the dense fallback's + # iterate-change stopping test. Each value is a maximum over the world. + r_b = float32(0.0) + r_p = float32(0.0) + r_d = float32(0.0) + r_c = float32(0.0) + + # Bilateral rows require v_aug = 0. + for jid in range(njc): + v_j = state_v_aug[vio + jid] + r_b = wp.max(r_b, wp.abs(v_j)) + + # Limits require lambda and v_aug in R+ with lambda * v_aug = 0. + for lid in range(nl): + lcio = vio + lcgo + lid + lambda_l = solution_lambdas[lcio] + v_l = state_v_aug[lcio] + r_p = wp.max(r_p, wp.abs(lambda_l - wp.max(0.0, lambda_l))) + r_d = wp.max(r_d, wp.abs(v_l - wp.max(0.0, v_l))) + r_c = wp.max(r_c, wp.abs(lambda_l * v_l)) + + # Contacts require lambda in K_mu, v_aug in its dual cone, and orthogonality. + for cid in range(nc): + ccio = vio + ccgo + 3 * cid + mu_c = problem_mu[cio + cid] + lambda_c = vec3f(solution_lambdas[ccio], solution_lambdas[ccio + 1], solution_lambdas[ccio + 2]) + v_c = vec3f(state_v_aug[ccio], state_v_aug[ccio + 1], state_v_aug[ccio + 2]) + lambda_proj = project_to_coulomb_cone(lambda_c, mu_c) + v_proj = project_to_coulomb_dual_cone(v_c, mu_c) + r_p = wp.max(r_p, wp.max(wp.abs(lambda_c - lambda_proj))) + r_d = wp.max(r_d, wp.max(wp.abs(v_c - v_proj))) + r_c = wp.max(r_c, wp.abs(wp.dot(lambda_c, v_c))) + + # Thus r_p and r_d are infinity-norm cone-projection distances, while r_c + # is the maximum absolute impulse-velocity inner product. + status.r_b = r_b + status.r_p = r_p + status.r_d = wp.max(r_d, r_b) + status.r_c = r_c + status.converged = int32(0) + if ncts == 0 or (r_b <= cfg.tolerance and r_p <= cfg.tolerance and r_d <= cfg.tolerance and r_c <= cfg.tolerance): + status.converged = int32(1) + solver_status[wid] = status + + +@wp.kernel +def _solve_dvi_pgs( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_vio: wp.array[int32], + problem_njc: wp.array[int32], + problem_nl: wp.array[int32], + problem_nc: wp.array[int32], + problem_lcgo: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + problem_D: wp.array[float32], + problem_v_f: wp.array[float32], + contact_block_inv: wp.array[mat33f], + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solver_status: wp.array[DVIStatus], + solution_lambdas: wp.array[float32], +): + wid = wp.tid() + + ncts = problem_dim[wid] + vio = problem_vio[wid] + mio = problem_mio[wid] + njc = problem_njc[wid] + nl = problem_nl[wid] + nc = problem_nc[wid] + lcgo = problem_lcgo[wid] + ccgo = problem_ccgo[wid] + cio = problem_cio[wid] + cfg = solver_config[wid] + + status = DVIStatus() + status.converged = int32(0) + status.iterations = int32(0) + status.r_p = float32(0.0) + status.r_d = float32(0.0) + status.r_c = float32(0.0) + status.r_b = float32(0.0) + + if ncts == 0: + status.converged = int32(1) + solver_status[wid] = status + return + + # This fallback stops on the unnormalized infinity norm of the impulse + # update. The residual fields below are provisional; solve() later replaces + # them with terminal cone-feasibility and complementarity diagnostics. + done = int32(0) + for iteration in range(cfg.max_iterations): + if done == 0: + max_step = float32(0.0) + max_velocity = float32(0.0) + max_complementarity = float32(0.0) + + # Equality constraints use scalar Gauss-Seidel updates. This keeps the + # solve on Kamino's Delassus system while avoiding a separate subsystem. + for i in range(njc): + v_i = _compute_row_velocity(ncts, mio, vio, i, problem_D, problem_v_f, solution_lambdas) + D_ii = wp.abs(problem_D[mio + ncts * i + i]) + cfg.regularization + FLOAT32_EPS + # Bilateral projection is the identity: lambda += -omega * B * v. + delta = -cfg.omega * v_i / D_ii + solution_lambdas[vio + i] += delta + max_step = wp.max(max_step, wp.abs(delta)) + max_velocity = wp.max(max_velocity, wp.abs(v_i)) + + for li in range(nl): + i = lcgo + li + v_i = _compute_row_velocity(ncts, mio, vio, i, problem_D, problem_v_f, solution_lambdas) + D_ii_raw = wp.abs(problem_D[mio + ncts * i + i]) + lambda_limit_old = solution_lambdas[vio + i] + lambda_limit_new = lambda_limit_old + if D_ii_raw > FLOAT32_EPS: + # Project lambda - omega * B * v onto the nonnegative ray. + lambda_limit_new = wp.max(0.0, lambda_limit_old - cfg.omega * v_i / (D_ii_raw + cfg.regularization)) + solution_lambdas[vio + i] = lambda_limit_new + max_step = wp.max(max_step, wp.abs(lambda_limit_new - lambda_limit_old)) + max_velocity = wp.max(max_velocity, wp.abs(wp.min(v_i, 0.0))) + max_complementarity = wp.max(max_complementarity, wp.abs(lambda_limit_new * v_i)) + + for cid in range(nc): + ccio = ccgo + 3 * cid + v_c = _contact_velocity_aug( + ncts, + mio, + vio, + ccgo, + cio, + cid, + problem_D, + problem_v_f, + solution_lambdas, + problem_mu, + ) + D_00 = wp.abs(problem_D[mio + ncts * (ccio + 0) + (ccio + 0)]) + D_11 = wp.abs(problem_D[mio + ncts * (ccio + 1) + (ccio + 1)]) + D_22 = wp.abs(problem_D[mio + ncts * (ccio + 2) + (ccio + 2)]) + lambda_contact_old = vec3f( + solution_lambdas[vio + ccio + 0], + solution_lambdas[vio + ccio + 1], + solution_lambdas[vio + ccio + 2], + ) + # Project the three-row impulse update onto K_mu. The block + # preconditioner retains normal-tangential coupling in D_cc. + if cfg.contact_block_preconditioner: + lambda_contact_projected = _project_contact_block_update( + lambda_contact_old, + v_c, + vec3f(D_00, D_11, D_22), + contact_block_inv[cio + cid], + cfg.regularization, + cfg.contact_jacobi_omega, + problem_mu[cio + cid], + ) + lambda_contact_new = lambda_contact_old + cfg.contact_jacobi_relaxation * ( + lambda_contact_projected - lambda_contact_old + ) + else: + lambda_contact_new = _project_contact_diagonal_update( + lambda_contact_old, + v_c, + _contact_trace_preconditioner(vec3f(D_00, D_11, D_22)), + cfg.regularization, + cfg.omega, + problem_mu[cio + cid], + ) + solution_lambdas[vio + ccio + 0] = lambda_contact_new.x + solution_lambdas[vio + ccio + 1] = lambda_contact_new.y + solution_lambdas[vio + ccio + 2] = lambda_contact_new.z + lambda_delta = lambda_contact_new - lambda_contact_old + max_step = wp.max(max_step, wp.max(wp.abs(lambda_delta))) + max_velocity = wp.max(max_velocity, wp.max(wp.abs(v_c))) + max_complementarity = wp.max(max_complementarity, wp.abs(wp.dot(lambda_contact_new, v_c))) + + status.iterations = iteration + int32(1) + status.r_p = max_step + status.r_d = max_velocity + status.r_c = max_complementarity + if max_step <= cfg.tolerance: + status.converged = int32(1) + done = int32(1) + + if done == 0: + status.converged = int32(0) + + solver_status[wid] = status + + +@wp.kernel +def _initialize_dvi_status( + # Inputs: + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solver_status: wp.array[DVIStatus], +): + wid = wp.tid() + cfg = solver_config[wid] + status = DVIStatus() + status.converged = int32(0) + status.iterations = cfg.contact_iterations + status.r_p = float32(0.0) + status.r_d = float32(0.0) + status.r_c = float32(0.0) + status.r_b = float32(0.0) + solver_status[wid] = status + + +@wp.kernel +def _set_dvi_direct_status_iterations( + # Inputs: + problem_nl: wp.array[int32], + problem_nc: wp.array[int32], + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solver_status: wp.array[DVIStatus], +): + wid = wp.tid() + cfg = solver_config[wid] + status = solver_status[wid] + if problem_nl[wid] == int32(0) and problem_nc[wid] == int32(0): + status.iterations = int32(1) + else: + status.iterations = cfg.block_iterations * cfg.contact_iterations + solver_status[wid] = status + + +@wp.kernel +def _set_dvi_bilateral_active_dim( + # Inputs: + problem_njc: wp.array[int32], + problem_nl: wp.array[int32], + problem_nc: wp.array[int32], + # Outputs: + bilateral_active_dim: wp.array[int32], +): + wid = wp.tid() + active_dim = int32(0) + if problem_nl[wid] > int32(0) or problem_nc[wid] > int32(0): + active_dim = problem_njc[wid] + bilateral_active_dim[wid] = active_dim + + +@wp.kernel +def _solve_dvi_limits_pgs( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_vio: wp.array[int32], + problem_nl: wp.array[int32], + problem_lcgo: wp.array[int32], + problem_D: wp.array[float32], + problem_v_f: wp.array[float32], + block_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solver_status: wp.array[DVIStatus], + solution_lambdas: wp.array[float32], +): + wid = wp.tid() + + ncts = problem_dim[wid] + vio = problem_vio[wid] + mio = problem_mio[wid] + nl = problem_nl[wid] + lcgo = problem_lcgo[wid] + cfg = solver_config[wid] + + if block_iteration >= cfg.block_iterations: + return + + status = DVIStatus() + status.converged = int32(0) + status.iterations = cfg.contact_iterations + status.r_p = float32(0.0) + status.r_d = float32(0.0) + status.r_c = float32(0.0) + status.r_b = float32(0.0) + + if ncts == 0 or nl == 0: + status.converged = int32(1) + solver_status[wid] = status + return + + # Limit sweeps use the same iterate-change stopping measure as the full + # dense fallback. Terminal DVI residuals are evaluated after all blocks. + done = int32(0) + for iteration in range(cfg.contact_iterations): + if done == 0: + max_step = float32(0.0) + max_velocity = float32(0.0) + max_complementarity = float32(0.0) + + for li in range(nl): + i = lcgo + li + v_i = _compute_row_velocity(ncts, mio, vio, i, problem_D, problem_v_f, solution_lambdas) + D_ii_raw = wp.abs(problem_D[mio + ncts * i + i]) + lambda_limit_old = solution_lambdas[vio + i] + lambda_limit_new = lambda_limit_old + if D_ii_raw > FLOAT32_EPS: + lambda_limit_new = wp.max(0.0, lambda_limit_old - cfg.omega * v_i / (D_ii_raw + cfg.regularization)) + solution_lambdas[vio + i] = lambda_limit_new + max_step = wp.max(max_step, wp.abs(lambda_limit_new - lambda_limit_old)) + max_velocity = wp.max(max_velocity, wp.abs(wp.min(v_i, 0.0))) + max_complementarity = wp.max(max_complementarity, wp.abs(lambda_limit_new * v_i)) + + status.iterations = iteration + int32(1) + status.r_p = max_step + status.r_d = max_velocity + status.r_c = max_complementarity + if max_step <= cfg.tolerance: + status.converged = int32(1) + done = int32(1) + + if done == 0: + status.converged = int32(0) + + solver_status[wid] = status + + +@wp.func_native(""" +#if defined(__CUDA_ARCH__) +__syncthreads(); +#endif +""") +def _sync_threads(): ... + + +@wp.func +def _contacts_share_dynamic_body(a: wp.vec2i, b: wp.vec2i) -> bool: + a0 = a[0] + a1 = a[1] + b0 = b[0] + b1 = b[1] + share = bool(False) + if a0 >= int32(0): + if a0 == b0 or a0 == b1: + share = bool(True) + if a1 >= int32(0): + if a1 == b0 or a1 == b1: + share = bool(True) + return share + + +@wp.kernel +def _color_dvi_contacts( + # Inputs: + problem_nc: wp.array[int32], + problem_cio: wp.array[int32], + contact_bid_AB: wp.array[wp.vec2i], + # Outputs: + contact_colors: wp.array[int32], + contact_num_colors: wp.array[int32], +): + wid = wp.tid() + + nc = problem_nc[wid] + cio = problem_cio[wid] + if nc == 0: + contact_num_colors[wid] = int32(0) + return + + # Contacts in one color share no dynamic body, so their Delassus + # cross-blocks vanish and their Gauss-Seidel updates may run concurrently. + num_colors = int32(0) + for cid in range(nc): + pair = contact_bid_AB[cio + cid] + color = int32(0) + found = int32(0) + while found == int32(0) and color < nc: + conflict = int32(0) + prev = int32(0) + while prev < cid: + if contact_colors[cio + prev] == color: + prev_pair = contact_bid_AB[cio + prev] + if _contacts_share_dynamic_body(pair, prev_pair): + conflict = int32(1) + prev = prev + int32(1) + + if conflict == int32(0): + found = int32(1) + else: + color = color + int32(1) + + contact_colors[cio + cid] = color + num_colors = wp.max(num_colors, color + int32(1)) + + contact_num_colors[wid] = num_colors + + +@wp.kernel +def _solve_dvi_contacts_colored_gs( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_vio: wp.array[int32], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + problem_D: wp.array[float32], + block_iteration: int32, + contact_block_inv: wp.array[mat33f], + contact_colors: wp.array[int32], + contact_num_colors: wp.array[int32], + solver_config: wp.array[DVIConfigStruct], + # Outputs: + state_v_aug: wp.array[float32], + solution_lambdas: wp.array[float32], +): + tid = wp.tid() + threads_per_world = int32(wp.block_dim()) + lane = tid % threads_per_world + wid = tid / threads_per_world + + nc = problem_nc[wid] + cio = problem_cio[wid] + if nc == 0: + return + + ncts = problem_dim[wid] + vio = problem_vio[wid] + mio = problem_mio[wid] + ccgo = problem_ccgo[wid] + cfg = solver_config[wid] + if block_iteration >= cfg.block_iterations: + return + + num_colors = contact_num_colors[wid] + # Colored contact updates execute a fixed number of sweeps. Convergence is + # evaluated only after the complete direct-bilateral block schedule. + iteration = int32(0) + while iteration < cfg.contact_iterations: + color = int32(0) + while color < num_colors: + cid = lane + while cid < nc: + if contact_colors[cio + cid] == color: + ccio = ccgo + int32(3) * cid + ccio_v = vio + ccio + mu_c = problem_mu[cio + cid] + + v_t0 = state_v_aug[ccio_v + 0] + v_t1 = state_v_aug[ccio_v + 1] + v_n = state_v_aug[ccio_v + 2] + mu_c * wp.sqrt(v_t0 * v_t0 + v_t1 * v_t1) + v_c = vec3f(v_t0, v_t1, v_n) + + D_00 = wp.abs(problem_D[mio + ncts * (ccio + 0) + (ccio + 0)]) + D_11 = wp.abs(problem_D[mio + ncts * (ccio + 1) + (ccio + 1)]) + D_22 = wp.abs(problem_D[mio + ncts * (ccio + 2) + (ccio + 2)]) + + lambda_old = vec3f( + solution_lambdas[ccio_v + 0], + solution_lambdas[ccio_v + 1], + solution_lambdas[ccio_v + 2], + ) + if cfg.contact_block_preconditioner: + lambda_projected = _project_contact_block_update( + lambda_old, + v_c, + vec3f(D_00, D_11, D_22), + contact_block_inv[cio + cid], + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + lambda_new = lambda_old + cfg.contact_jacobi_relaxation * (lambda_projected - lambda_old) + else: + lambda_new = _project_contact_diagonal_update( + lambda_old, + v_c, + _contact_trace_preconditioner(vec3f(D_00, D_11, D_22)), + cfg.regularization, + cfg.omega, + mu_c, + ) + + delta = lambda_new - lambda_old + solution_lambdas[ccio_v + 0] = lambda_new.x + solution_lambdas[ccio_v + 1] = lambda_new.y + solution_lambdas[ccio_v + 2] = lambda_new.z + + # Only contact velocities are read before all velocities are rebuilt. + row = ccgo + contact_end = ccgo + int32(3) * nc + while row < contact_end: + row_mio = mio + ncts * row + dv = ( + problem_D[row_mio + ccio + 0] * delta.x + + problem_D[row_mio + ccio + 1] * delta.y + + problem_D[row_mio + ccio + 2] * delta.z + ) + wp.atomic_add(state_v_aug, vio + row, dv) + row = row + int32(1) + + cid = cid + threads_per_world + + _sync_threads() + color = color + int32(1) + iteration = iteration + int32(1) + + +@wp.kernel +def _compute_dvi_contact_jacobi_delta( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_vio: wp.array[int32], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + problem_D: wp.array[float32], + block_iteration: int32, + contact_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + contact_block_inv: wp.array[mat33f], + state_v_aug: wp.array[float32], + # Outputs: + solution_lambdas: wp.array[float32], + state_scratch: wp.array[float32], +): + wid, cid = wp.tid() + + nc = problem_nc[wid] + if cid >= nc: + return + + ncts = problem_dim[wid] + vio = problem_vio[wid] + mio = problem_mio[wid] + ccgo = problem_ccgo[wid] + ccio = ccgo + int32(3) * cid + ccio_v = vio + ccio + mu_c = problem_mu[problem_cio[wid] + cid] + cfg = solver_config[wid] + if block_iteration >= cfg.block_iterations or contact_iteration >= cfg.contact_iterations: + return + + # All contacts read the same velocity snapshot. Store impulse deltas so the + # companion kernel can apply their coupled Delassus effect simultaneously. + v_t0 = state_v_aug[ccio_v + 0] + v_t1 = state_v_aug[ccio_v + 1] + v_n = state_v_aug[ccio_v + 2] + mu_c * wp.sqrt(v_t0 * v_t0 + v_t1 * v_t1) + v_c = vec3f(v_t0, v_t1, v_n) + D_00 = wp.abs(problem_D[mio + ncts * (ccio + 0) + (ccio + 0)]) + D_11 = wp.abs(problem_D[mio + ncts * (ccio + 1) + (ccio + 1)]) + D_22 = wp.abs(problem_D[mio + ncts * (ccio + 2) + (ccio + 2)]) + + lambda_old = vec3f( + solution_lambdas[ccio_v + 0], + solution_lambdas[ccio_v + 1], + solution_lambdas[ccio_v + 2], + ) + if cfg.contact_block_preconditioner: + lambda_projected = _project_contact_block_update( + lambda_old, + v_c, + vec3f(D_00, D_11, D_22), + contact_block_inv[problem_cio[wid] + cid], + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + else: + lambda_projected = _project_contact_diagonal_update( + lambda_old, + v_c, + _contact_trace_preconditioner(vec3f(D_00, D_11, D_22)), + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + lambda_new = lambda_old + cfg.contact_jacobi_relaxation * (lambda_projected - lambda_old) + + delta = lambda_new - lambda_old + solution_lambdas[ccio_v + 0] = lambda_new.x + solution_lambdas[ccio_v + 1] = lambda_new.y + solution_lambdas[ccio_v + 2] = lambda_new.z + state_scratch[ccio_v + 0] = delta.x + state_scratch[ccio_v + 1] = delta.y + state_scratch[ccio_v + 2] = delta.z + + +@wp.kernel +def _apply_dvi_contact_jacobi_delta( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_vio: wp.array[int32], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_D: wp.array[float32], + block_iteration: int32, + contact_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + state_scratch: wp.array[float32], + # Outputs: + state_v_aug: wp.array[float32], +): + wid, tid = wp.tid() + + cfg = solver_config[wid] + if block_iteration >= cfg.block_iterations or contact_iteration >= cfg.contact_iterations: + return + + nc = problem_nc[wid] + if tid >= int32(3) * nc: + return + + ncts = problem_dim[wid] + vio = problem_vio[wid] + mio = problem_mio[wid] + ccgo = problem_ccgo[wid] + row = ccgo + tid + row_mio = mio + ncts * row + + # Accumulate D_cc * delta_lambda for the Jacobi sweep; no contact observes + # another contact's new impulse until every delta has been formed. + dv = float32(0.0) + for cid in range(nc): + ccio = ccgo + int32(3) * cid + dv += problem_D[row_mio + ccio + 0] * state_scratch[vio + ccio + 0] + dv += problem_D[row_mio + ccio + 1] * state_scratch[vio + ccio + 1] + dv += problem_D[row_mio + ccio + 2] * state_scratch[vio + ccio + 2] + + state_v_aug[vio + row] += dv + + +@wp.kernel +def _compute_dvi_contact_velocities( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_vio: wp.array[int32], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_D: wp.array[float32], + problem_v_f: wp.array[float32], + solution_lambdas: wp.array[float32], + # Outputs: + state_v_aug: wp.array[float32], +): + wid, tid = wp.tid() + + if tid >= int32(3) * problem_nc[wid]: + return + + ncts = problem_dim[wid] + mio = problem_mio[wid] + vio = problem_vio[wid] + row = problem_ccgo[wid] + tid + state_v_aug[vio + row] = _compute_row_velocity(ncts, mio, vio, row, problem_D, problem_v_f, solution_lambdas) + + +@wp.kernel +def _compute_dvi_solution_vectors( + # Inputs: + problem_dim: wp.array[int32], + problem_mio: wp.array[int32], + problem_vio: wp.array[int32], + problem_D: wp.array[float32], + problem_v_f: wp.array[float32], + # Outputs: + state_s: wp.array[float32], + state_v_aug: wp.array[float32], + solution_lambdas: wp.array[float32], + solution_v_plus: wp.array[float32], +): + wid, tid = wp.tid() + + ncts = problem_dim[wid] + if tid >= ncts: + return + + mio = problem_mio[wid] + vio = problem_vio[wid] + # Recover the physical post-event velocity v_plus = D * lambda + v_f. + # De Saxce augmentation is stored separately for cone residual evaluation. + v_i = _compute_row_velocity(ncts, mio, vio, tid, problem_D, problem_v_f, solution_lambdas) + solution_v_plus[vio + tid] = v_i + state_v_aug[vio + tid] = v_i + state_s[vio + tid] = 0.0 + + +@wp.kernel +def _compute_dvi_desaxce_corrections( + # Inputs: + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_vio: wp.array[int32], + problem_mu: wp.array[float32], + # Outputs: + state_s: wp.array[float32], + state_v_aug: wp.array[float32], + solution_v_plus: wp.array[float32], +): + wid, cid = wp.tid() + + nc = problem_nc[wid] + if cid >= nc: + return + + vio = problem_vio[wid] + ccgo = problem_ccgo[wid] + ccio = ccgo + 3 * cid + vt0 = solution_v_plus[vio + ccio] + vt1 = solution_v_plus[vio + ccio + 1] + # s = [0, 0, mu * ||v_t||] maps physical contact velocity to the dual-cone + # variable v_aug = v_plus + s used by the DVI contact conditions. + s_n = problem_mu[problem_cio[wid] + cid] * wp.sqrt(vt0 * vt0 + vt1 * vt1) + state_s[vio + ccio + 2] = s_n + state_v_aug[vio + ccio + 2] = solution_v_plus[vio + ccio + 2] + s_n + + +@wp.kernel +def _unprecondition_dvi_solution( + # Inputs: + problem_dim: wp.array[int32], + problem_vio: wp.array[int32], + problem_P: wp.array[float32], + # Outputs: + state_s: wp.array[float32], + state_v_aug: wp.array[float32], + solution_lambdas: wp.array[float32], + solution_v_plus: wp.array[float32], +): + wid, tid = wp.tid() + + ncts = problem_dim[wid] + if tid >= ncts: + return + + vio = problem_vio[wid] + v_i = vio + tid + P_i = problem_P[v_i] + # The solver uses D_hat = P * D * P: impulses map with P, while + # constraint-space velocities and De Saxce terms map with P^-1. + solution_lambdas[v_i] = P_i * solution_lambdas[v_i] + solution_v_plus[v_i] = solution_v_plus[v_i] / P_i + state_v_aug[v_i] = state_v_aug[v_i] / P_i + state_s[v_i] = state_s[v_i] / P_i diff --git a/newton/_src/solvers/kamino/_src/solvers/dvi/projections.py b/newton/_src/solvers/kamino/_src/solvers/dvi/projections.py new file mode 100644 index 0000000000..5538f283b9 --- /dev/null +++ b/newton/_src/solvers/kamino/_src/solvers/dvi/projections.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Shared contact projection functions for dense and sparse DVI kernels.""" + +import warp as wp + +from ...core.math import FLOAT32_EPS +from ..padmm.math import project_to_coulomb_cone + +float32 = wp.float32 +mat33f = wp.mat33f +vec3f = wp.vec3f + + +@wp.func +def project_contact_diagonal_update( + lambda_old: vec3f, + v_c: vec3f, + D_diag: vec3f, + regularization: float32, + omega: float32, + mu: float32, +) -> vec3f: + """Apply a diagonally preconditioned contact projection. + + Computes ``lambda_next = project_K(lambda - omega * B * v_aug)``. + """ + lambda_arg = lambda_old + if D_diag.x > FLOAT32_EPS: + lambda_arg.x = lambda_old.x - omega * v_c.x / (D_diag.x + regularization) + if D_diag.y > FLOAT32_EPS: + lambda_arg.y = lambda_old.y - omega * v_c.y / (D_diag.y + regularization) + if D_diag.z > FLOAT32_EPS: + lambda_arg.z = lambda_old.z - omega * v_c.z / (D_diag.z + regularization) + return project_to_coulomb_cone(lambda_arg, mu) + + +@wp.func +def project_contact_block_update( + lambda_old: vec3f, + v_c: vec3f, + D_diag: vec3f, + D_block_inv: mat33f, + regularization: float32, + omega: float32, + mu: float32, +) -> vec3f: + """Apply a block-preconditioned contact projection. + + Computes ``lambda_next = project_K(lambda - omega * B * v_aug)``. + """ + inv_diag_norm = wp.abs(D_block_inv[0, 0]) + wp.abs(D_block_inv[1, 1]) + wp.abs(D_block_inv[2, 2]) + if inv_diag_norm > FLOAT32_EPS: + return project_to_coulomb_cone(lambda_old - omega * (D_block_inv * v_c), mu) + return project_contact_diagonal_update(lambda_old, v_c, D_diag, regularization, omega, mu) + + +@wp.func +def contact_trace_preconditioner(D_diag: vec3f) -> vec3f: + """Return the isotropic contact preconditioner derived from block trace.""" + D_eff = (D_diag.x + D_diag.y + D_diag.z) / float32(3.0) + return vec3f(D_eff, D_eff, D_eff) diff --git a/newton/_src/solvers/kamino/_src/solvers/dvi/solver.py b/newton/_src/solvers/kamino/_src/solvers/dvi/solver.py new file mode 100644 index 0000000000..6fcf589854 --- /dev/null +++ b/newton/_src/solvers/kamino/_src/solvers/dvi/solver.py @@ -0,0 +1,875 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Projected DVI solver for Kamino dual forward-dynamics problems.""" + +from __future__ import annotations + +import warp as wp + +from ....config import DVISolverConfig +from ...core.data import DataKamino +from ...core.model import ModelKamino +from ...core.size import SizeKamino +from ...core.types import to_warp_int32_array +from ...dynamics.dual import DualProblem +from ...geometry.contacts import ContactsKamino +from ...kinematics.jacobians import SparseSystemJacobians +from ...kinematics.limits import LimitsKamino +from ...linalg import DenseLinearOperatorData, DenseSquareMultiLinearInfo, LLTBlockedRCMSolver, LLTBlockedSolver +from ..common import ( + WarmStartMode, + apply_dual_preconditioner_to_solution, + warmstart_contact_constraints, + warmstart_joint_constraints, + warmstart_limit_constraints, +) +from .kernels import ( + _apply_dvi_contact_jacobi_delta, + _build_bilateral_rhs, + _color_dvi_contacts, + _compute_dvi_contact_block_inverse, + _compute_dvi_contact_jacobi_delta, + _compute_dvi_contact_velocities, + _compute_dvi_desaxce_corrections, + _compute_dvi_solution_vectors, + _compute_dvi_status_residuals, + _copy_bilateral_block, + _initialize_dvi_status, + _reset_dvi_solver_data, + _reset_dvi_status, + _scatter_bilateral_solution, + _set_dvi_bilateral_active_dim, + _set_dvi_direct_status_iterations, + _solve_dvi_contacts_colored_gs, + _solve_dvi_limits_pgs, + _solve_dvi_pgs, + _unprecondition_dvi_solution, +) +from .sparse import SparseDVIPath +from .types import DVIConfigStruct, DVIData, convert_config_to_struct + +wp.set_module_options({"enable_backward": False}) + +float32 = wp.float32 + + +class DVISolver: + """Solve Kamino dual problems with projected DVI iterations. + + For Kamino's dual system ``v_plus = D * lambda + v_f``, bilateral rows + enforce zero velocity, limit rows enforce nonnegative complementarity, + and contact rows enforce Coulomb-cone complementarity after the De Saxce + velocity correction. + + Bilateral constraints are solved as a direct block when available, while + limits and frictional contacts use projected Gauss-Seidel, Jacobi, or + graph-colored updates. Dense and matrix-free sparse problems share the + same solution, warm-start, status, and diagnostics contract. + """ + + Config = DVISolverConfig + + def __init__( + self, + model: ModelKamino | None = None, + data: DataKamino | None = None, + limits: LimitsKamino | None = None, + contacts: ContactsKamino | None = None, + jacobians: SparseSystemJacobians | None = None, + problem: DualProblem | None = None, + config: list[DVISolver.Config] | DVISolver.Config | None = None, + warmstart: WarmStartMode = WarmStartMode.NONE, + collect_info: bool = False, + ): + """Initialize a DVI solver and optionally allocate it for a model. + + Args: + model: Model that determines solver allocation sizes. + data: Model data used by sparse DVI operator products. + limits: Limit topology used by sparse DVI updates. + contacts: Contact topology used for graph-colored contact updates. + jacobians: Sparse constraint Jacobians used by sparse DVI updates. + problem: Optional sparse problem used to precompute topology outside + the first simulation step. + config: One DVI config or one config per world. + warmstart: Source used to initialize constraint impulses. + collect_info: Whether to retain terminal per-world diagnostics. + """ + self._config: list[DVISolver.Config] = [] + self._warmstart: WarmStartMode = WarmStartMode.NONE + self._collect_info: bool = False + self._size: SizeKamino | None = None + self._data: DVIData | None = None + self._bilateral_solver: LLTBlockedSolver | LLTBlockedRCMSolver | None = None + self._max_block_iterations: int = 1 + self._max_contact_iterations: int = 1 + self._max_iterations: int = 1 + self._bilateral_solve_after_block: tuple[bool, ...] = () + self._has_contact_block_preconditioner: bool = False + self._has_unilateral_constraints: bool = False + self._contact_bid_AB: wp.array[wp.vec2i] | None = None + self._sparse_path: SparseDVIPath | None = None + self._all_worlds_mask: wp.array[wp.bool] | None = None + self._device: wp.DeviceLike = None + + if model is not None: + self.finalize( + model=model, + data=data, + limits=limits, + contacts=contacts, + jacobians=jacobians, + problem=problem, + config=config, + warmstart=warmstart, + collect_info=collect_info, + ) + + @property + def config(self) -> list[DVISolver.Config]: + """Host-side per-world DVI configs.""" + return self._config + + @property + def size(self) -> SizeKamino: + """Model size cache.""" + return self._size + + @property + def data(self) -> DVIData: + """Solver data arrays.""" + if self._data is None: + raise RuntimeError("Solver data has not been allocated yet. Call `finalize()` first.") + return self._data + + @property + def device(self) -> wp.DeviceLike: + """Device on which solver data is allocated.""" + return self._device + + @property + def all_worlds_mask(self) -> wp.array[wp.bool]: + """Boolean mask selecting every world for sparse operator products.""" + return self._all_worlds_mask + + def finalize( + self, + model: ModelKamino, + data: DataKamino | None = None, + limits: LimitsKamino | None = None, + contacts: ContactsKamino | None = None, + jacobians: SparseSystemJacobians | None = None, + problem: DualProblem | None = None, + config: list[DVISolver.Config] | DVISolver.Config | None = None, + warmstart: WarmStartMode = WarmStartMode.NONE, + collect_info: bool = False, + ): + """Allocate solver data and precompute model-dependent topology. + + Args: + model: Model that determines solver allocation sizes. + data: Model data used by sparse DVI operator products. + limits: Limit topology used by sparse DVI updates. + contacts: Contact topology used for graph-colored contact updates. + jacobians: Sparse constraint Jacobians used by sparse DVI updates. + problem: Optional sparse problem used to precompute topology. + config: One DVI config or one config per world. + warmstart: Source used to initialize constraint impulses. + collect_info: Whether to retain terminal per-world diagnostics. + """ + if model is None or not isinstance(model, ModelKamino): + raise ValueError("A model of type `ModelKamino` must be provided.") + + self._size = model.size + self._device = model.device + self._config = self._check_config(model, config) + self._warmstart = warmstart + self._collect_info = collect_info + self._max_iterations = max(c.max_iterations for c in self._config) + self._max_block_iterations = max(c.block_iterations for c in self._config) + self._max_contact_iterations = max(c.contact_iterations for c in self._config) + self._bilateral_solve_after_block = self._make_bilateral_solve_schedule(self._config) + self._has_contact_block_preconditioner = any(c.contact_block_preconditioner for c in self._config) + self._has_unilateral_constraints = self._size.max_of_max_limits > 0 or self._size.max_of_max_contacts > 0 + self._data = DVIData(size=self._size, collect_info=self._collect_info, device=self._device) + self._all_worlds_mask = wp.ones(shape=(self._size.num_worlds,), dtype=wp.bool, device=self._device) + self._allocate_bilateral_solver(model) + self._sparse_path = SparseDVIPath( + device=self._device, + size=self._size, + data=self._data, + model=model, + model_data=data, + limits=limits, + contacts=contacts, + jacobians=jacobians, + bilateral_solver=self._bilateral_solver, + max_iterations=self._max_iterations, + max_block_iterations=self._max_block_iterations, + max_contact_iterations=self._max_contact_iterations, + has_contact_block_preconditioner=self._has_contact_block_preconditioner, + has_unilateral_constraints=self._has_unilateral_constraints, + all_worlds_mask=self._all_worlds_mask, + should_solve_bilateral_after_block=self._should_solve_bilateral_after_block, + ) + self.set_contacts(contacts) + if problem is not None and problem.sparse: + self._sparse_path.prepare(problem) + + configs = [convert_config_to_struct(c) for c in self._config] + with wp.ScopedDevice(self._device): + self._data.config = wp.array(configs, dtype=DVIConfigStruct) + + def _make_bilateral_solve_schedule(self, configs: list[DVISolver.Config]) -> tuple[bool, ...]: + """Return host-side repeated bilateral solve points for direct-block DVI.""" + return tuple( + any(next_block < c.block_iterations and next_block % c.bilateral_solve_period == 0 for c in configs) + for next_block in range(1, self._max_block_iterations) + ) + + def _should_solve_bilateral_after_block(self, block_iteration: int) -> bool: + """Whether the direct bilateral block should be re-solved after this block.""" + if block_iteration < 0 or block_iteration >= len(self._bilateral_solve_after_block): + return False + return self._bilateral_solve_after_block[block_iteration] + + def _allocate_bilateral_solver(self, model: ModelKamino): + """Allocate the reduced dense operator used for bilateral DVI solves.""" + self._bilateral_solver = None + self._data.bilateral_operator = None + if model.size.sum_of_num_joint_cts == 0: + return + + joint_cts_per_world = model.info.num_joint_cts.numpy().astype(int).tolist() + if any(njc <= 0 for njc in joint_cts_per_world): + return + + mat_sizes = [njc * njc for njc in joint_cts_per_world] + mat_offsets = [0] + for size in mat_sizes[:-1]: + mat_offsets.append(mat_offsets[-1] + size) + + operator = DenseLinearOperatorData() + operator.info = DenseSquareMultiLinearInfo() + operator.info.assign( + maxdim=model.info.num_joint_cts, + dim=model.info.num_joint_cts, + mio=to_warp_int32_array(mat_offsets, device=self._device), + vio=model.info.joint_cts_offset, + dtype=float32, + device=self._device, + ) + operator.mat = wp.zeros(shape=(operator.info.total_mat_size,), dtype=float32, device=self._device) + self._data.bilateral_operator = operator + first_config = self._config[0] + if any( + config.bilateral_solver_type != first_config.bilateral_solver_type + or config.bilateral_solver_kwargs != first_config.bilateral_solver_kwargs + for config in self._config[1:] + ): + raise ValueError("All worlds must use the same DVI bilateral solver configuration.") + + solver_type = first_config.bilateral_solver_type + kwargs = dict(first_config.bilateral_solver_kwargs) + if solver_type == "LLTB": + # A larger factorization tile reduces panel count, while the + # single-RHS solve benefits from more threads on its smaller tile. + kwargs.setdefault("factorize_block_size", 64) + kwargs.setdefault("solve_block_dim", 256) + solver_class = LLTBlockedSolver + else: + solver_class = LLTBlockedRCMSolver + self._bilateral_solver = solver_class(operator=operator, device=self._device, **kwargs) + + @staticmethod + def _check_config( + model: ModelKamino | None = None, config: list[DVISolver.Config] | DVISolver.Config | None = None + ) -> list[DVISolver.Config]: + if config is None: + config = [DVISolver.Config()] * (model.info.num_worlds if model else 1) + elif isinstance(config, DVISolver.Config): + config = [config] * (model.info.num_worlds if model else 1) + elif isinstance(config, list): + if model is not None and len(config) != model.info.num_worlds: + raise ValueError(f"Expected {model.info.num_worlds} configs, got {len(config)}") + if not all(isinstance(c, DVISolver.Config) for c in config): + raise TypeError("All configs must be instances of DVISolver.Config") + else: + raise TypeError(f"Expected a single object or list of `DVISolver.Config`, got {type(config)}") + return config + + def set_contacts(self, contacts: ContactsKamino | None): + """Cache contact topology for graph-colored contact solves.""" + if contacts is not None and contacts.model_max_contacts_host > 0: + self._contact_bid_AB = contacts.bid_AB + else: + self._contact_bid_AB = None + + def reset(self, problem: DualProblem | None = None, world_mask: wp.array[wp.bool] | None = None): + """Reset scratch state and cached solution data.""" + self._data.state.reset() + if self._data.info is not None: + self._data.info.zero() + if world_mask is None: + self._data.solution.zero() + else: + if problem is None: + raise ValueError("A `DualProblem` instance must be provided when a world mask is used.") + wp.launch( + kernel=_reset_dvi_solver_data, + dim=(self._size.num_worlds, self._size.max_of_max_total_cts), + inputs=[ + world_mask, + problem.data.vio, + problem.data.maxdim, + self._data.solution.lambdas, + self._data.solution.v_plus, + ], + device=self.device, + ) + + def coldstart(self): + """Prepare a cold-start solve.""" + self._data.state.reset() + self._data.solution.zero() + + def warmstart( + self, + problem: DualProblem, + model: ModelKamino, + data: DataKamino, + limits: LimitsKamino | None = None, + contacts: ContactsKamino | None = None, + ): + """Prepare a warm-start solve.""" + self._data.state.reset() + self.set_contacts(contacts) + + match self._warmstart: + case WarmStartMode.NONE: + self._data.solution.zero() + case WarmStartMode.INTERNAL: + self._warmstart_from_solution(problem) + case WarmStartMode.CONTAINERS: + self._warmstart_from_containers(problem, model, data, limits, contacts) + case _: + raise ValueError(f"Invalid warmstart mode: {self._warmstart}") + + def solve(self, problem: DualProblem): + """Solve the cone-complementarity problem defined by ``problem``. + + Kamino supplies the constraint-space system + + ``v_plus = D * lambda + v_f``, + + where ``D`` is Kamino's represented Delassus operator derived from + ``J * M^-1 * J^T``, ``lambda`` contains joint, limit, and contact + impulses, and ``v_f`` contains unconstrained motion, stabilization, + and restitution. The bilateral equation ``D * lambda = -v_f`` is the + standalone ``N * lambda = b`` formulation with ``b = -v_f``. + + Rows are ordered as bilateral joints, unilateral limits, and contact + triplets ``[t0, t1, n]``. Contact rows use + ``v_aug = v_plus + [0, 0, mu * norm(v_t)]``. + + The DVI solution satisfies zero ``v_aug`` on bilateral rows, + nonnegative complementarity on limit rows, and Coulomb-cone + complementarity between contact impulses and augmented velocities. DVI + exploits this difference by partitioning the system into bilateral + impulses ``lambda_b`` and unilateral limit/contact impulses + ``lambda_u``. When the bilateral block is available, it is factored and + solved directly: + + ``D_bb * lambda_b = -(v_f,b + D_bu * lambda_u)``. + + The unilateral block is updated iteratively with projection onto the + nonnegative and Coulomb cones. Alternating these updates retains the + ``D_bu`` and ``D_ub`` coupling while using a solver suited to each + constraint class. Repeating the alternation for ``block_iterations`` + drives ``lambda_b`` and ``lambda_u`` toward a mutually consistent + solution; a single block without a bilateral re-solve reduces to a + one-directional solve where the joints never see the final contact and + limit impulses. The fallback path instead applies projected + Gauss-Seidel to all rows. + + This differs from Kamino's PADMM backend, which places all constraint + rows in one proximal-ADMM iteration: it solves a regularized full + Delassus system for the unconstrained primal update, then projects the + unilateral components. DVI uses no ADMM penalty or auxiliary-variable + iteration; its primary split is direct bilateral versus projected + iterative unilateral solves. Dense and sparse DVI paths implement the + same split with different Delassus representations. + + Args: + problem: Unified Kamino dual problem to solve. + """ + wp.launch( + kernel=_reset_dvi_status, + dim=self._size.num_worlds, + inputs=[self._data.status], + device=self.device, + ) + + if problem.sparse: + if self._sparse_path is None: + raise RuntimeError("Sparse DVI path has not been allocated. Call `finalize()` first.") + if self._sparse_path.bilateral_nzb_pairs is None: + self._sparse_path.prepare(problem) + # Apply projected iterations through matrix-free products + # D * lambda = J * (M^-1 * (J^T * lambda)) + R * lambda. + self._sparse_path.solve(problem) + elif self._has_contact_block_preconditioner and self._size.max_of_max_contacts > 0: + # For each contact c, form B_c = (D_cc + regularization * I)^-1. + wp.launch( + kernel=_compute_dvi_contact_block_inverse, + dim=(self._size.num_worlds, self._size.max_of_max_contacts), + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.nc, + problem.data.ccgo, + problem.data.cio, + problem.data.D, + self._data.config, + self._data.state.contact_block_inv, + ], + device=self.device, + ) + + if not problem.sparse: + if self._bilateral_solver is not None and self._data.bilateral_operator is not None: + self._solve_with_bilateral_direct_block(problem) + else: + # Solve all rows together with projected Gauss-Seidel: + # lambda_next = projection(lambda - omega * B * v_aug). + wp.launch( + kernel=_solve_dvi_pgs, + dim=self._size.num_worlds, + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.vio, + problem.data.njc, + problem.data.nl, + problem.data.nc, + problem.data.lcgo, + problem.data.ccgo, + problem.data.cio, + problem.data.mu, + problem.data.D, + problem.data.v_f, + self._data.state.contact_block_inv, + self._data.config, + self._data.status, + self._data.solution.lambdas, + ], + device=self.device, + ) + + # Evaluate the physical post-event velocity v_plus = D * lambda + v_f. + wp.launch( + kernel=_compute_dvi_solution_vectors, + dim=(self._size.num_worlds, self._size.max_of_max_total_cts), + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.vio, + problem.data.D, + problem.data.v_f, + self._data.state.s, + self._data.state.v_aug, + self._data.solution.lambdas, + self._data.solution.v_plus, + ], + device=self.device, + ) + + if self._size.max_of_max_contacts > 0: + # Map physical contact velocity to the dual-cone variable + # v_aug = v_plus + [0, 0, mu * norm(v_t)]. + wp.launch( + kernel=_compute_dvi_desaxce_corrections, + dim=(self._size.num_worlds, self._size.max_of_max_contacts), + inputs=[ + problem.data.nc, + problem.data.ccgo, + problem.data.cio, + problem.data.vio, + problem.data.mu, + self._data.state.s, + self._data.state.v_aug, + self._data.solution.v_plus, + ], + device=self.device, + ) + + # Classify the final iterate using all DVI conditions. This replaces + # provisional iterate-change convergence from the dense fallback; + # direct and sparse paths reach this check after fixed iteration counts. + wp.launch( + kernel=_compute_dvi_status_residuals, + dim=self._size.num_worlds, + inputs=[ + problem.data.dim, + problem.data.vio, + problem.data.njc, + problem.data.nl, + problem.data.nc, + problem.data.lcgo, + problem.data.ccgo, + problem.data.cio, + problem.data.mu, + self._data.config, + self._data.state.v_aug, + self._data.solution.lambdas, + self._data.status, + ], + device=self.device, + ) + + if self._collect_info: + wp.copy(self._data.info.status, self._data.status) + + wp.launch( + kernel=_unprecondition_dvi_solution, + dim=(self._size.num_worlds, self._size.max_of_max_total_cts), + inputs=[ + problem.data.dim, + problem.data.vio, + problem.data.P, + self._data.state.s, + self._data.state.v_aug, + self._data.solution.lambdas, + self._data.solution.v_plus, + ], + device=self.device, + ) + + def _solve_bilateral_block(self, problem: DualProblem, active_dim: wp.array[wp.int32] | None = None): + """Solve ``D_bb * lambda_b = -(v_f,b + D_bu * lambda_u)``.""" + operator = self._data.bilateral_operator + state = self._data.state + wp.launch( + kernel=_build_bilateral_rhs, + dim=(self._size.num_worlds, self._size.max_of_num_joint_cts), + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.vio, + problem.data.njc, + problem.data.D, + problem.data.v_f, + operator.info.vio, + state.bilateral_preconditioner, + self._data.solution.lambdas, + state.bilateral_rhs, + ], + device=self.device, + ) + full_dim = operator.info.dim + if active_dim is not None: + operator.info.dim = active_dim + try: + self._bilateral_solver.solve(b=state.bilateral_rhs, x=state.bilateral_solution) + finally: + operator.info.dim = full_dim + wp.launch( + kernel=_scatter_bilateral_solution, + dim=(self._size.num_worlds, self._size.max_of_num_joint_cts), + inputs=[ + problem.data.vio, + problem.data.njc, + operator.info.vio, + state.bilateral_preconditioner, + state.bilateral_solution, + self._data.solution.lambdas, + ], + device=self.device, + ) + + def _factor_bilateral_block(self, problem: DualProblem): + """Extract, symmetrically scale, and factor the bilateral block ``D_bb``.""" + operator = self._data.bilateral_operator + operator.info.dim = operator.info.maxdim + wp.launch( + kernel=_copy_bilateral_block, + dim=(self._size.num_worlds, self._size.max_of_num_joint_cts * self._size.max_of_num_joint_cts), + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.njc, + problem.data.D, + operator.info.mio, + operator.info.vio, + operator.mat, + self._data.state.bilateral_preconditioner, + ], + device=self.device, + ) + self._bilateral_solver.compute(A=operator.mat) + + def _solve_with_bilateral_direct_block(self, problem: DualProblem): + """Alternate a direct bilateral solve with projected unilateral updates. + + With unilateral impulses fixed, the direct solve satisfies + ``D_bb * lambda_b = -(v_f,b + D_bu * lambda_u)``. Repeating these + updates preserves bilateral-unilateral coupling. Between direct solves, + limits and contacts apply projected updates using the unilateral + residual ``D_ub * lambda_b + D_uu * lambda_u + v_f,u``. As the block + count grows the two impulse sets converge to a mutually consistent + solution; one block without a bilateral re-solve corresponds to a + one-directional joint-then-contact solve. + """ + self._factor_bilateral_block(problem) + self._solve_bilateral_block(problem) + if not self._has_unilateral_constraints: + return + + wp.launch( + kernel=_initialize_dvi_status, + dim=self._size.num_worlds, + inputs=[ + self._data.config, + self._data.status, + ], + device=self.device, + ) + + wp.launch( + kernel=_set_dvi_bilateral_active_dim, + dim=self._size.num_worlds, + inputs=[ + problem.data.njc, + problem.data.nl, + problem.data.nc, + self._data.state.bilateral_active_dim, + ], + device=self.device, + ) + + use_colored_contacts = ( + self._size.max_of_max_contacts > 0 and self.device.is_cuda and self._contact_bid_AB is not None + ) + if use_colored_contacts: + wp.launch( + kernel=_color_dvi_contacts, + dim=self._size.num_worlds, + inputs=[ + problem.data.nc, + problem.data.cio, + self._contact_bid_AB, + self._data.state.contact_colors, + self._data.state.contact_num_colors, + ], + device=self.device, + ) + + for block_iteration in range(self._max_block_iterations): + if self._size.max_of_max_limits > 0: + wp.launch( + kernel=_solve_dvi_limits_pgs, + dim=self._size.num_worlds, + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.vio, + problem.data.nl, + problem.data.lcgo, + problem.data.D, + problem.data.v_f, + block_iteration, + self._data.config, + self._data.status, + self._data.solution.lambdas, + ], + device=self.device, + ) + + if self._size.max_of_max_contacts > 0: + wp.launch( + kernel=_compute_dvi_contact_velocities, + dim=(self._size.num_worlds, 3 * self._size.max_of_max_contacts), + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.vio, + problem.data.nc, + problem.data.ccgo, + problem.data.D, + problem.data.v_f, + self._data.solution.lambdas, + self._data.state.v_aug, + ], + device=self.device, + ) + + if use_colored_contacts: + wp.launch( + kernel=_solve_dvi_contacts_colored_gs, + dim=self._size.num_worlds * 64, + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.vio, + problem.data.nc, + problem.data.ccgo, + problem.data.cio, + problem.data.mu, + problem.data.D, + block_iteration, + self._data.state.contact_block_inv, + self._data.state.contact_colors, + self._data.state.contact_num_colors, + self._data.config, + self._data.state.v_aug, + self._data.solution.lambdas, + ], + device=self.device, + block_dim=64, + ) + else: + for contact_iteration in range(self._max_contact_iterations): + wp.launch( + kernel=_compute_dvi_contact_jacobi_delta, + dim=(self._size.num_worlds, self._size.max_of_max_contacts), + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.vio, + problem.data.nc, + problem.data.ccgo, + problem.data.cio, + problem.data.mu, + problem.data.D, + block_iteration, + contact_iteration, + self._data.config, + self._data.state.contact_block_inv, + self._data.state.v_aug, + self._data.solution.lambdas, + self._data.state.scratch, + ], + device=self.device, + ) + wp.launch( + kernel=_apply_dvi_contact_jacobi_delta, + dim=(self._size.num_worlds, 3 * self._size.max_of_max_contacts), + inputs=[ + problem.data.dim, + problem.data.mio, + problem.data.vio, + problem.data.nc, + problem.data.ccgo, + problem.data.D, + block_iteration, + contact_iteration, + self._data.config, + self._data.state.scratch, + self._data.state.v_aug, + ], + device=self.device, + ) + + if self._should_solve_bilateral_after_block(block_iteration): + self._solve_bilateral_block(problem, active_dim=self._data.state.bilateral_active_dim) + + self._solve_bilateral_block(problem, active_dim=self._data.state.bilateral_active_dim) + + wp.launch( + kernel=_set_dvi_direct_status_iterations, + dim=self._size.num_worlds, + inputs=[ + problem.data.nl, + problem.data.nc, + self._data.config, + self._data.status, + ], + device=self.device, + ) + + def _warmstart_from_solution(self, problem: DualProblem): + wp.launch( + kernel=apply_dual_preconditioner_to_solution, + dim=(self._size.num_worlds, self._size.max_of_max_total_cts), + inputs=[ + problem.data.dim, + problem.data.vio, + problem.data.P, + self._data.solution.lambdas, + self._data.solution.v_plus, + ], + device=self.device, + ) + + def _warmstart_from_containers( + self, + problem: DualProblem, + model: ModelKamino, + data: DataKamino, + limits: LimitsKamino | None = None, + contacts: ContactsKamino | None = None, + ): + self._data.solution.zero() + if model.size.sum_of_num_joints > 0: + wp.launch( + kernel=warmstart_joint_constraints, + dim=model.size.sum_of_num_joints, + inputs=[ + model.time.dt, + model.joints.wid, + model.joints.num_dynamic_cts, + model.joints.num_kinematic_cts, + model.joints.dynamic_cts_offset_joint_cts, + model.joints.kinematic_cts_offset_joint_cts, + model.joints.dynamic_cts_offset_total_cts, + model.joints.kinematic_cts_offset_total_cts, + data.joints.lambda_j, + problem.data.P, + self._data.solution.lambdas, + self._data.solution.lambdas, + self._data.solution.v_plus, + ], + device=self.device, + ) + if limits is not None and limits.model_max_limits_host > 0: + wp.launch( + kernel=warmstart_limit_constraints, + dim=limits.model_max_limits_host, + inputs=[ + model.time.dt, + model.info.total_cts_offset, + data.info.limit_cts_group_offset, + limits.model_active_limits, + limits.wid, + limits.lid, + limits.reaction, + limits.velocity, + problem.data.P, + self._data.solution.lambdas, + self._data.solution.lambdas, + self._data.solution.v_plus, + ], + device=self.device, + ) + if contacts is not None and contacts.model_max_contacts_host > 0: + wp.launch( + kernel=warmstart_contact_constraints, + dim=contacts.model_max_contacts_host, + inputs=[ + model.time.dt, + model.info.total_cts_offset, + data.info.contact_cts_group_offset, + contacts.model_active_contacts, + contacts.wid, + contacts.cid, + contacts.material, + contacts.reaction, + contacts.velocity, + problem.data.P, + self._data.solution.lambdas, + self._data.solution.lambdas, + self._data.solution.v_plus, + ], + device=self.device, + ) diff --git a/newton/_src/solvers/kamino/_src/solvers/dvi/sparse.py b/newton/_src/solvers/kamino/_src/solvers/dvi/sparse.py new file mode 100644 index 0000000000..464a4239c4 --- /dev/null +++ b/newton/_src/solvers/kamino/_src/solvers/dvi/sparse.py @@ -0,0 +1,658 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Sparse DVI solve path for Kamino dual systems.""" + +from __future__ import annotations + +import warp as wp + +from ...core.data import DataKamino +from ...core.model import ModelKamino +from ...dynamics.delassus import BlockSparseMatrixFreeDelassusOperator +from ...dynamics.dual import DualProblem +from ...geometry.contacts import ContactsKamino +from ...kinematics.jacobians import SparseSystemJacobians +from ...kinematics.limits import LimitsKamino +from . import sparse_kernels +from .kernels import ( + _initialize_dvi_status, + _scatter_bilateral_solution, + _set_dvi_bilateral_active_dim, + _set_dvi_direct_status_iterations, +) +from .sparse_kernels import ( + _build_sparse_bilateral_block, + _build_sparse_bilateral_rhs, + _compute_dvi_sparse_solution_vectors, + _set_dvi_sparse_status_iterations, + _set_sparse_bilateral_diagonal, + _solve_dvi_sparse_contacts_offset_update, + _solve_dvi_sparse_jacobi_update, + _solve_dvi_sparse_limits_offset_update, + _solve_dvi_sparse_unilateral_jacobi_update, + _solve_dvi_sparse_unilateral_offset_update, + _sparse_delassus_gemv_rows, + _zero_bilateral_lambdas, +) + +wp.set_module_options({"enable_backward": False}) + +int32 = wp.int32 + + +_SPARSE_DELASSUS_ROWS_JOINTS = 0 +_SPARSE_DELASSUS_ROWS_UNILATERAL = 1 + + +class SparseDVIPath: + """Own workspace and operations for the sparse Kamino DVI solve path.""" + + def __init__( + self, + device: wp.DeviceLike, + size, + data, + model: ModelKamino, + model_data: DataKamino | None, + limits: LimitsKamino | None, + contacts: ContactsKamino | None, + jacobians: SparseSystemJacobians | None, + bilateral_solver, + max_iterations: int, + max_block_iterations: int, + max_contact_iterations: int, + has_contact_block_preconditioner: bool, + has_unilateral_constraints: bool, + all_worlds_mask: wp.array[wp.bool], + should_solve_bilateral_after_block, + ): + """Initialize the sparse-path workspace references.""" + self.device = device + self.size = size + self.data = data + self.model = model + self.model_data = model_data + self.limits = limits + self.contacts = contacts + self.jacobians = jacobians + self.body_space = wp.empty(shape=size.sum_of_num_body_dofs, dtype=wp.float32, device=device) + self.bilateral_solver = bilateral_solver + self.max_iterations = max_iterations + self.max_block_iterations = max_block_iterations + self.max_contact_iterations = max_contact_iterations + self.has_contact_block_preconditioner = has_contact_block_preconditioner + self.has_unilateral_constraints = has_unilateral_constraints + self.all_worlds_mask = all_worlds_mask + self.should_solve_bilateral_after_block = should_solve_bilateral_after_block + self.bilateral_nzb_pairs: ( + tuple[ + wp.array[wp.int32], + wp.array[wp.int32], + wp.array[wp.int32], + wp.array[wp.int32], + wp.array[wp.int32], + wp.array[wp.int32], + ] + | None + ) = None + + def prepare(self, problem: DualProblem) -> None: + """Precompute host-derived sparse topology before the first solve.""" + _get_sparse_delassus(problem) + if self.model_data is None or self.jacobians is None: + raise RuntimeError("Sparse DVI requires model data and sparse Jacobians.") + if self.bilateral_solver is not None and self.data.bilateral_operator is not None: + _build_sparse_bilateral_pairs(self, problem) + + def solve(self, problem: DualProblem) -> None: + """Solve a sparse Kamino DVI problem without materializing dense Delassus.""" + if self.has_contact_block_preconditioner and self.size.max_of_max_contacts > 0: + _compute_sparse_contact_block_inverse(self, problem) + + if self.bilateral_solver is not None and self.data.bilateral_operator is not None: + _solve_sparse_with_bilateral_direct_block(self, problem) + else: + _solve_sparse_jacobi(self, problem) + + +def _get_sparse_delassus(problem: DualProblem) -> BlockSparseMatrixFreeDelassusOperator: + delassus = problem.delassus + if not isinstance(delassus, BlockSparseMatrixFreeDelassusOperator): + raise TypeError("Sparse DVI requires a `BlockSparseMatrixFreeDelassusOperator`.") + return delassus + + +def _solve_sparse_jacobi(path: SparseDVIPath, problem: DualProblem) -> None: + """Apply fixed Jacobi sweeps to the unified sparse DVI problem. + + Each sweep evaluates ``v_aug = D * lambda + v_f + s`` matrix-free, then + applies ``lambda_next = projection(lambda - omega * B * v_aug)``. + """ + state = path.data.state + problem.delassus.diagonal(state.scratch) + + for iteration in range(path.max_iterations): + problem.delassus.matvec( + x=path.data.solution.lambdas, + y=state.v_aug, + world_mask=path.all_worlds_mask, + ) + wp.launch( + kernel=_solve_dvi_sparse_jacobi_update, + dim=(path.size.num_worlds, path.size.max_of_max_total_cts), + inputs=[ + problem.data.dim, + problem.data.vio, + problem.data.njc, + problem.data.nl, + problem.data.nc, + problem.data.lcgo, + problem.data.ccgo, + problem.data.cio, + problem.data.mu, + state.scratch, + problem.data.P, + problem.data.v_f, + state.v_aug, + state.contact_block_inv, + iteration, + path.data.config, + path.data.solution.lambdas, + ], + device=path.device, + ) + + problem.delassus.matvec( + x=path.data.solution.lambdas, + y=state.v_aug, + world_mask=path.all_worlds_mask, + ) + wp.launch( + kernel=_compute_dvi_sparse_solution_vectors, + dim=(path.size.num_worlds, path.size.max_of_max_total_cts), + inputs=[ + problem.data.dim, + problem.data.vio, + problem.data.v_f, + state.s, + state.v_aug, + path.data.solution.v_plus, + ], + device=path.device, + ) + wp.launch( + kernel=_set_dvi_sparse_status_iterations, + dim=path.size.num_worlds, + inputs=[ + problem.data.dim, + path.data.config, + path.data.status, + ], + device=path.device, + ) + + +def _compute_sparse_solution_vectors(path: SparseDVIPath, problem: DualProblem) -> None: + state = path.data.state + problem.delassus.matvec( + x=path.data.solution.lambdas, + y=state.v_aug, + world_mask=path.all_worlds_mask, + ) + wp.launch( + kernel=_compute_dvi_sparse_solution_vectors, + dim=(path.size.num_worlds, path.size.max_of_max_total_cts), + inputs=[ + problem.data.dim, + problem.data.vio, + problem.data.v_f, + state.s, + state.v_aug, + path.data.solution.v_plus, + ], + device=path.device, + ) + + +def _sparse_delassus_matvec_rows_path(path: SparseDVIPath, problem: DualProblem, row_kind: int) -> None: + delassus = _get_sparse_delassus(problem) + state = path.data.state + regularization = delassus.regularization + body_space = path.body_space + bsm = delassus.bsm + if bsm is None: + raise RuntimeError("Sparse DVI row products require initialized Delassus sparse operators.") + + # Evaluate selected rows of D * lambda = J * M^-1 * J^T * lambda + R * lambda + # without materializing the Delassus matrix. + delassus.apply_jacobian_transpose(path.data.solution.lambdas, body_space, path.all_worlds_mask) + state.v_aug.zero_() + wp.launch( + kernel=_sparse_delassus_gemv_rows, + dim=(bsm.num_matrices, bsm.max_of_num_nzb), + inputs=[ + bsm.dims, + bsm.num_nzb, + bsm.nzb_start, + bsm.nzb_coords, + bsm.nzb_values, + bsm.row_start, + bsm.col_start, + problem.data.dim, + problem.data.njc, + row_kind, + regularization, + body_space, + state.v_aug, + path.data.solution.lambdas, + path.all_worlds_mask, + ], + device=path.device, + ) + + +def _sparse_delassus_matvec_rows(solver, problem: DualProblem, row_kind: int) -> None: + """Compatibility wrapper for sparse Delassus row products.""" + if solver._sparse_path is None: + raise RuntimeError("Sparse DVI path has not been allocated. Call `finalize()` first.") + _sparse_delassus_matvec_rows_path(solver._sparse_path, problem, row_kind) + + +def _sparse_delassus_update_unilateral_rows( + path: SparseDVIPath, + problem: DualProblem, + block_iteration: int, + contact_iteration: int, +) -> None: + if _sparse_delassus_update_unilateral_offsets(path, problem, block_iteration, contact_iteration): + return + + state = path.data.state + _sparse_delassus_matvec_rows_path(path, problem, _SPARSE_DELASSUS_ROWS_UNILATERAL) + wp.launch( + kernel=_solve_dvi_sparse_unilateral_jacobi_update, + dim=(path.size.num_worlds, path.size.max_of_max_total_cts), + inputs=[ + problem.data.dim, + problem.data.vio, + problem.data.njc, + problem.data.nl, + problem.data.nc, + problem.data.lcgo, + problem.data.ccgo, + problem.data.cio, + problem.data.mu, + state.scratch, + problem.data.P, + problem.data.v_f, + state.v_aug, + state.contact_block_inv, + block_iteration, + contact_iteration, + path.data.config, + path.data.solution.lambdas, + ], + device=path.device, + ) + + +def _sparse_delassus_update_unilateral_offsets( + path: SparseDVIPath, + problem: DualProblem, + block_iteration: int, + contact_iteration: int, +) -> bool: + delassus = _get_sparse_delassus(problem) + state = path.data.state + regularization = delassus.regularization + body_space = path.body_space + bsm = delassus.bsm + jacobians = path.jacobians + limits = path.limits + contacts = path.contacts + limit_offsets = jacobians.limit_constraint_nzb_offsets + contact_offsets = jacobians.contact_constraint_nzb_offsets + has_limits = limits is not None and limits.model_max_limits_host > 0 + has_contacts = contacts is not None and contacts.model_max_contacts_host > 0 + + if not (has_limits or has_contacts): + return False + if bsm is None: + raise RuntimeError("Sparse DVI offset updates require initialized Delassus sparse operators.") + + delassus.apply_jacobian_transpose(path.data.solution.lambdas, body_space, path.all_worlds_mask) + + if has_limits and has_contacts: + # Fuse the two independent sweeps (disjoint lambda outputs, shared + # body_space) into one launch to remove a per-iteration kernel launch. + limits_capacity = limits.model_max_limits_host + wp.launch( + kernel=_solve_dvi_sparse_unilateral_offset_update, + dim=limits_capacity + contacts.model_max_contacts_host, + inputs=[ + limits_capacity, + bsm.num_nzb, + bsm.nzb_start, + bsm.nzb_coords, + bsm.nzb_values, + bsm.row_start, + bsm.col_start, + limits.model_active_limits, + limits.wid, + limits.lid, + limit_offsets, + problem.data.nl, + problem.data.lcgo, + contacts.model_active_contacts, + contacts.wid, + contacts.cid, + contact_offsets, + problem.data.nc, + problem.data.ccgo, + problem.data.cio, + problem.data.mu, + state.contact_block_inv, + problem.data.vio, + state.scratch, + problem.data.P, + problem.data.v_f, + regularization, + body_space, + block_iteration, + contact_iteration, + path.data.config, + path.data.solution.lambdas, + ], + device=path.device, + ) + return True + + if has_limits: + wp.launch( + kernel=_solve_dvi_sparse_limits_offset_update, + dim=limits.model_max_limits_host, + inputs=[ + bsm.num_nzb, + bsm.nzb_start, + bsm.nzb_coords, + bsm.nzb_values, + bsm.row_start, + bsm.col_start, + limits.model_active_limits, + limits.wid, + limits.lid, + limit_offsets, + problem.data.vio, + problem.data.nl, + problem.data.lcgo, + state.scratch, + problem.data.P, + problem.data.v_f, + regularization, + body_space, + block_iteration, + contact_iteration, + path.data.config, + path.data.solution.lambdas, + ], + device=path.device, + ) + + if has_contacts: + wp.launch( + kernel=_solve_dvi_sparse_contacts_offset_update, + dim=contacts.model_max_contacts_host, + inputs=[ + bsm.num_nzb, + bsm.nzb_start, + bsm.nzb_coords, + bsm.nzb_values, + bsm.row_start, + bsm.col_start, + contacts.model_active_contacts, + contacts.wid, + contacts.cid, + contact_offsets, + problem.data.vio, + problem.data.nc, + problem.data.ccgo, + problem.data.cio, + problem.data.mu, + state.scratch, + problem.data.P, + problem.data.v_f, + regularization, + body_space, + state.contact_block_inv, + block_iteration, + contact_iteration, + path.data.config, + path.data.solution.lambdas, + ], + device=path.device, + ) + + return True + + +def _compute_sparse_contact_block_inverse(path: SparseDVIPath, problem: DualProblem) -> None: + jacobian = problem.delassus.constraint_jacobian + wp.launch( + kernel=sparse_kernels._compute_sparse_contact_block_inverse, + dim=(path.size.num_worlds, path.size.max_of_max_contacts), + inputs=[ + path.model.info.bodies_offset, + path.model.bodies.inv_m_i, + path.model_data.bodies.inv_I_i, + jacobian.nzb_start, + jacobian.num_nzb, + jacobian.nzb_coords, + jacobian.nzb_values, + problem.data.nc, + problem.data.ccgo, + problem.data.cio, + problem.data.vio, + problem.data.P, + path.data.config, + jacobian.max_of_num_nzb, + path.data.state.contact_block_inv, + ], + device=path.device, + ) + + +def _factor_sparse_bilateral_block(path: SparseDVIPath, problem: DualProblem) -> None: + operator = path.data.bilateral_operator + state = path.data.state + operator.info.dim = operator.info.maxdim + operator.mat.zero_() + state.bilateral_preconditioner.zero_() + problem.delassus.diagonal(state.scratch) + + jacobian = problem.delassus.constraint_jacobian + if path.bilateral_nzb_pairs is None: + raise RuntimeError("Sparse DVI topology is not prepared. Call `SparseDVIPath.prepare()` before solving.") + wp.launch( + kernel=_set_sparse_bilateral_diagonal, + dim=(path.size.num_worlds, path.size.max_of_num_joint_cts), + inputs=[ + problem.data.njc, + problem.data.vio, + operator.info.mio, + operator.info.vio, + state.scratch, + operator.mat, + state.bilateral_preconditioner, + ], + device=path.device, + ) + pair_wid, pair_row, pair_col, pair_bid, pair_i, pair_j = path.bilateral_nzb_pairs + if pair_wid.size > 0: + wp.launch( + kernel=_build_sparse_bilateral_block, + dim=pair_wid.size, + inputs=[ + path.model.bodies.inv_m_i, + path.model_data.bodies.inv_I_i, + pair_wid, + pair_row, + pair_col, + pair_bid, + pair_i, + pair_j, + jacobian.nzb_values, + problem.data.njc, + operator.info.mio, + operator.info.vio, + state.bilateral_preconditioner, + operator.mat, + ], + device=path.device, + ) + path.bilateral_solver.compute(A=operator.mat) + + +def _build_sparse_bilateral_pairs(path: SparseDVIPath, problem: DualProblem) -> None: + """Cache joint Jacobian block pairs that contribute to the bilateral matrix.""" + jacobian = problem.delassus.constraint_jacobian + counts = path.jacobians.joint_constraint_nzb_count.numpy().tolist() + starts = jacobian.nzb_start.numpy().tolist() + coords = jacobian.nzb_coords.numpy() + joint_counts = problem.data.njc.numpy().tolist() + body_offsets = path.model.info.bodies_offset.numpy().tolist() + + pair_wid: list[int] = [] + pair_row: list[int] = [] + pair_col: list[int] = [] + pair_bid: list[int] = [] + pair_i: list[int] = [] + pair_j: list[int] = [] + for wid, count in enumerate(counts): + start = starts[wid] + njc = joint_counts[wid] + for local_i in range(count): + nzb_i = start + local_i + row = int(coords[nzb_i, 0]) + body_col = int(coords[nzb_i, 1]) + if row >= njc: + continue + for local_j in range(count): + nzb_j = start + local_j + col = int(coords[nzb_j, 0]) + if row < col < njc and body_col == int(coords[nzb_j, 1]): + pair_wid.append(wid) + pair_row.append(row) + pair_col.append(col) + pair_bid.append(body_offsets[wid] + body_col // 6) + pair_i.append(nzb_i) + pair_j.append(nzb_j) + + path.bilateral_nzb_pairs = tuple( + wp.array(values, dtype=int32, device=path.device) + for values in (pair_wid, pair_row, pair_col, pair_bid, pair_i, pair_j) + ) + + +def _solve_sparse_bilateral_block( + path: SparseDVIPath, problem: DualProblem, active_dim: wp.array[int32] | None = None +) -> None: + operator = path.data.bilateral_operator + state = path.data.state + wp.launch( + kernel=_zero_bilateral_lambdas, + dim=(path.size.num_worlds, path.size.max_of_num_joint_cts), + inputs=[ + problem.data.njc, + problem.data.vio, + path.data.solution.lambdas, + ], + device=path.device, + ) + _sparse_delassus_matvec_rows_path(path, problem, _SPARSE_DELASSUS_ROWS_JOINTS) + wp.launch( + kernel=_build_sparse_bilateral_rhs, + dim=(path.size.num_worlds, path.size.max_of_num_joint_cts), + inputs=[ + problem.data.vio, + problem.data.njc, + problem.data.v_f, + state.v_aug, + operator.info.vio, + state.bilateral_preconditioner, + state.bilateral_rhs, + ], + device=path.device, + ) + full_dim = operator.info.dim + if active_dim is not None: + operator.info.dim = active_dim + try: + path.bilateral_solver.solve(b=state.bilateral_rhs, x=state.bilateral_solution) + finally: + operator.info.dim = full_dim + wp.launch( + kernel=_scatter_bilateral_solution, + dim=(path.size.num_worlds, path.size.max_of_num_joint_cts), + inputs=[ + problem.data.vio, + problem.data.njc, + operator.info.vio, + state.bilateral_preconditioner, + state.bilateral_solution, + path.data.solution.lambdas, + ], + device=path.device, + ) + + +def _solve_sparse_with_bilateral_direct_block(path: SparseDVIPath, problem: DualProblem) -> None: + """Alternate a direct ``D_bb`` solve with projected sparse unilateral sweeps.""" + state = path.data.state + _factor_sparse_bilateral_block(path, problem) + _solve_sparse_bilateral_block(path, problem) + if not path.has_unilateral_constraints: + _compute_sparse_solution_vectors(path, problem) + return + + wp.launch( + kernel=_initialize_dvi_status, + dim=path.size.num_worlds, + inputs=[ + path.data.config, + path.data.status, + ], + device=path.device, + ) + wp.launch( + kernel=_set_dvi_bilateral_active_dim, + dim=path.size.num_worlds, + inputs=[ + problem.data.njc, + problem.data.nl, + problem.data.nc, + state.bilateral_active_dim, + ], + device=path.device, + ) + + for block_iteration in range(path.max_block_iterations): + for contact_iteration in range(path.max_contact_iterations): + _sparse_delassus_update_unilateral_rows(path, problem, block_iteration, contact_iteration) + + if path.should_solve_bilateral_after_block(block_iteration): + _solve_sparse_bilateral_block(path, problem, active_dim=state.bilateral_active_dim) + + _solve_sparse_bilateral_block(path, problem, active_dim=state.bilateral_active_dim) + wp.launch( + kernel=_set_dvi_direct_status_iterations, + dim=path.size.num_worlds, + inputs=[ + problem.data.nl, + problem.data.nc, + path.data.config, + path.data.status, + ], + device=path.device, + ) + _compute_sparse_solution_vectors(path, problem) diff --git a/newton/_src/solvers/kamino/_src/solvers/dvi/sparse_kernels.py b/newton/_src/solvers/kamino/_src/solvers/dvi/sparse_kernels.py new file mode 100644 index 0000000000..39ade741b8 --- /dev/null +++ b/newton/_src/solvers/kamino/_src/solvers/dvi/sparse_kernels.py @@ -0,0 +1,993 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Sparse Warp kernels for the Kamino DVI solver.""" + +from __future__ import annotations + +import warp as wp + +from ...core.math import FLOAT32_EPS +from ...core.types import vec6f +from .projections import ( + contact_trace_preconditioner as _contact_trace_preconditioner, +) +from .projections import ( + project_contact_block_update as _project_contact_block_update, +) +from .projections import ( + project_contact_diagonal_update as _project_contact_diagonal_update, +) +from .types import DVIConfigStruct, DVIStatus + +wp.set_module_options({"enable_backward": False}) + +float32 = wp.float32 +int32 = wp.int32 +mat33f = wp.mat33f +vec3f = wp.vec3f + + +@wp.kernel +def _zero_bilateral_lambdas( + # Inputs: + problem_njc: wp.array[int32], + problem_vio: wp.array[int32], + # Outputs: + solution_lambdas: wp.array[float32], +): + wid, row = wp.tid() + + njc = problem_njc[wid] + if row >= njc: + return + + solution_lambdas[problem_vio[wid] + row] = 0.0 + + +@wp.kernel +def _build_sparse_bilateral_rhs( + # Inputs: + problem_vio: wp.array[int32], + problem_njc: wp.array[int32], + problem_v_f: wp.array[float32], + state_v_aug: wp.array[float32], + bilateral_vio: wp.array[int32], + bilateral_P: wp.array[float32], + # Outputs: + bilateral_rhs: wp.array[float32], +): + wid, row = wp.tid() + + njc = problem_njc[wid] + if row >= njc: + return + + pvio = problem_vio[wid] + bvio = bilateral_vio[wid] + rhs = -(state_v_aug[pvio + row] + problem_v_f[pvio + row]) + bilateral_rhs[bvio + row] = bilateral_P[bvio + row] * rhs + + +@wp.kernel +def _sparse_delassus_gemv_rows( + # Matrix data: + dims: wp.array2d[int32], + num_nzb: wp.array[int32], + nzb_start: wp.array[int32], + nzb_coords: wp.array2d[int32], + nzb_values: wp.array[vec6f], + row_start: wp.array[int32], + col_start: wp.array[int32], + # Row ranges: + problem_dim: wp.array[int32], + problem_njc: wp.array[int32], + row_kind: int32, + # Regularization: + eta: wp.array[float32], + # Vectors: + body_space: wp.array[float32], + y: wp.array[float32], + lambdas: wp.array[float32], + # Mask: + world_mask: wp.array[bool], +): + wid, block_idx = wp.tid() + + if not world_mask[wid]: + return + + dim = problem_dim[wid] + njc = problem_njc[wid] + + if block_idx < dim: + row = block_idx + row_active = row < njc + if row_kind == int32(1): + row_active = row >= njc + if row_active: + vec_idx = row_start[wid] + row + wp.atomic_add(y, vec_idx, eta[vec_idx] * lambdas[vec_idx]) + + if block_idx >= num_nzb[wid]: + return + + global_block_idx = nzb_start[wid] + block_idx + block_coord = nzb_coords[global_block_idx] + row = block_coord[0] + if row < 0 or row >= dim: + return + + row_active = row < njc + if row_kind == int32(1): + row_active = row >= njc + if not row_active: + return + + # The body-space input already contains M^-1 * J^T * lambda. Accumulate + # selected rows of J times that vector; eta * lambda supplies R * lambda. + block = nzb_values[global_block_idx] + x_idx_base = col_start[wid] + block_coord[1] + acc = float32(0.0) + for j in range(6): + acc += block[j] * body_space[x_idx_base + j] + + wp.atomic_add(y, row_start[wid] + row, acc) + + +@wp.func +def _apply_limit_offset_update( + limit_id: int32, + # Matrix data: + num_nzb: wp.array[int32], + nzb_start: wp.array[int32], + nzb_coords: wp.array2d[int32], + nzb_values: wp.array[vec6f], + row_start: wp.array[int32], + col_start: wp.array[int32], + # Active limits: + limits_model_active: wp.array[int32], + limits_wid: wp.array[int32], + limits_lid: wp.array[int32], + limits_nzb_offsets: wp.array[int32], + # Problem data: + problem_vio: wp.array[int32], + problem_nl: wp.array[int32], + problem_lcgo: wp.array[int32], + problem_diag: wp.array[float32], + problem_P: wp.array[float32], + problem_v_f: wp.array[float32], + eta: wp.array[float32], + body_space: wp.array[float32], + block_iteration: int32, + contact_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solution_lambdas: wp.array[float32], +): + if limit_id >= limits_model_active[0]: + return + + wid = limits_wid[limit_id] + cfg = solver_config[wid] + if block_iteration >= cfg.block_iterations or contact_iteration >= cfg.contact_iterations: + return + + lid = limits_lid[limit_id] + if lid >= problem_nl[wid]: + return + + row = problem_lcgo[wid] + lid + vec_idx = problem_vio[wid] + row + eta_idx = row_start[wid] + row + value = eta[eta_idx] * solution_lambdas[vec_idx] + matrix_end = nzb_start[wid] + num_nzb[wid] + nzb_offset = limits_nzb_offsets[limit_id] + for k in range(2): + nzb_idx = nzb_offset + k + if nzb_idx < matrix_end: + block_coord = nzb_coords[nzb_idx] + if block_coord[0] == row: + block = nzb_values[nzb_idx] + x_idx_base = col_start[wid] + block_coord[1] + for j in range(6): + value += block[j] * body_space[x_idx_base + j] + + v_i = value + problem_v_f[vec_idx] + P_i = problem_P[vec_idx] + D_ii_raw = wp.abs(problem_diag[vec_idx]) * P_i * P_i + D_ii = D_ii_raw + cfg.regularization + FLOAT32_EPS + + # Each active limit row references at most two body blocks. The resulting + # matrix-free velocity is followed by projection onto the nonnegative ray. + lambda_limit_old = solution_lambdas[vec_idx] + lambda_limit_new = lambda_limit_old + if D_ii_raw > FLOAT32_EPS: + lambda_limit_new = wp.max(0.0, lambda_limit_old - cfg.omega * v_i / D_ii) + solution_lambdas[vec_idx] = lambda_limit_new + + +@wp.kernel +def _solve_dvi_sparse_limits_offset_update( + # Matrix data: + num_nzb: wp.array[int32], + nzb_start: wp.array[int32], + nzb_coords: wp.array2d[int32], + nzb_values: wp.array[vec6f], + row_start: wp.array[int32], + col_start: wp.array[int32], + # Active limits: + limits_model_active: wp.array[int32], + limits_wid: wp.array[int32], + limits_lid: wp.array[int32], + limits_nzb_offsets: wp.array[int32], + # Problem data: + problem_vio: wp.array[int32], + problem_nl: wp.array[int32], + problem_lcgo: wp.array[int32], + problem_diag: wp.array[float32], + problem_P: wp.array[float32], + problem_v_f: wp.array[float32], + eta: wp.array[float32], + body_space: wp.array[float32], + block_iteration: int32, + contact_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solution_lambdas: wp.array[float32], +): + _apply_limit_offset_update( + wp.tid(), + num_nzb, + nzb_start, + nzb_coords, + nzb_values, + row_start, + col_start, + limits_model_active, + limits_wid, + limits_lid, + limits_nzb_offsets, + problem_vio, + problem_nl, + problem_lcgo, + problem_diag, + problem_P, + problem_v_f, + eta, + body_space, + block_iteration, + contact_iteration, + solver_config, + solution_lambdas, + ) + + +@wp.func +def _apply_contact_offset_update( + contact_id: int32, + # Matrix data: + num_nzb: wp.array[int32], + nzb_start: wp.array[int32], + nzb_coords: wp.array2d[int32], + nzb_values: wp.array[vec6f], + row_start: wp.array[int32], + col_start: wp.array[int32], + # Active contacts: + contacts_model_active: wp.array[int32], + contacts_wid: wp.array[int32], + contacts_cid: wp.array[int32], + contacts_nzb_offsets: wp.array[int32], + # Problem data: + problem_vio: wp.array[int32], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + problem_diag: wp.array[float32], + problem_P: wp.array[float32], + problem_v_f: wp.array[float32], + eta: wp.array[float32], + body_space: wp.array[float32], + contact_block_inv: wp.array[mat33f], + block_iteration: int32, + contact_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solution_lambdas: wp.array[float32], +): + if contact_id >= contacts_model_active[0]: + return + + wid = contacts_wid[contact_id] + cfg = solver_config[wid] + if block_iteration >= cfg.block_iterations or contact_iteration >= cfg.contact_iterations: + return + + cid = contacts_cid[contact_id] + nc = problem_nc[wid] + if cid >= nc: + return + + ccio = problem_ccgo[wid] + int32(3) * cid + row_0 = ccio + int32(0) + row_1 = ccio + int32(1) + row_2 = ccio + int32(2) + vio = problem_vio[wid] + ccio_v = vio + ccio + row_offset = row_start[wid] + + value_0 = eta[row_offset + row_0] * solution_lambdas[ccio_v + 0] + value_1 = eta[row_offset + row_1] * solution_lambdas[ccio_v + 1] + value_2 = eta[row_offset + row_2] * solution_lambdas[ccio_v + 2] + + matrix_end = nzb_start[wid] + num_nzb[wid] + nzb_offset = contacts_nzb_offsets[contact_id] + for k in range(3): + nzb_idx = nzb_offset + k + block = nzb_values[nzb_idx] + x_idx_base = col_start[wid] + nzb_coords[nzb_idx, 1] + acc = float32(0.0) + for j in range(6): + acc += block[j] * body_space[x_idx_base + j] + if k == 0: + value_0 += acc + elif k == 1: + value_1 += acc + else: + value_2 += acc + + second_body_offset = nzb_offset + 3 + if second_body_offset < matrix_end and nzb_coords[second_body_offset, 0] == row_0: + for k in range(3): + nzb_idx = second_body_offset + k + block = nzb_values[nzb_idx] + x_idx_base = col_start[wid] + nzb_coords[nzb_idx, 1] + acc = float32(0.0) + for j in range(6): + acc += block[j] * body_space[x_idx_base + j] + if k == 0: + value_0 += acc + elif k == 1: + value_1 += acc + else: + value_2 += acc + + mu_c = problem_mu[problem_cio[wid] + cid] + v_t0 = value_0 + problem_v_f[ccio_v + 0] + v_t1 = value_1 + problem_v_f[ccio_v + 1] + v_n = value_2 + problem_v_f[ccio_v + 2] + mu_c * wp.sqrt(v_t0 * v_t0 + v_t1 * v_t1) + v_c = vec3f(v_t0, v_t1, v_n) + + P_0 = problem_P[ccio_v + 0] + P_1 = problem_P[ccio_v + 1] + P_2 = problem_P[ccio_v + 2] + D_00 = wp.abs(problem_diag[ccio_v + 0]) * P_0 * P_0 + D_11 = wp.abs(problem_diag[ccio_v + 1]) * P_1 * P_1 + D_22 = wp.abs(problem_diag[ccio_v + 2]) * P_2 * P_2 + + # Contact topology offsets select the one or two body blocks contributing + # to this [t0, t1, n] row triplet before its Coulomb-cone projection. + lambda_contact_old = vec3f( + solution_lambdas[ccio_v + 0], + solution_lambdas[ccio_v + 1], + solution_lambdas[ccio_v + 2], + ) + D_diag = _contact_trace_preconditioner(vec3f(D_00, D_11, D_22)) + if cfg.contact_block_preconditioner: + lambda_projected = _project_contact_block_update( + lambda_contact_old, + v_c, + D_diag, + contact_block_inv[problem_cio[wid] + cid], + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + else: + lambda_projected = _project_contact_diagonal_update( + lambda_contact_old, + v_c, + D_diag, + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + lambda_contact_new = lambda_contact_old + cfg.contact_jacobi_relaxation * (lambda_projected - lambda_contact_old) + + solution_lambdas[ccio_v + 0] = lambda_contact_new.x + solution_lambdas[ccio_v + 1] = lambda_contact_new.y + solution_lambdas[ccio_v + 2] = lambda_contact_new.z + + +@wp.kernel +def _solve_dvi_sparse_contacts_offset_update( + # Matrix data: + num_nzb: wp.array[int32], + nzb_start: wp.array[int32], + nzb_coords: wp.array2d[int32], + nzb_values: wp.array[vec6f], + row_start: wp.array[int32], + col_start: wp.array[int32], + # Active contacts: + contacts_model_active: wp.array[int32], + contacts_wid: wp.array[int32], + contacts_cid: wp.array[int32], + contacts_nzb_offsets: wp.array[int32], + # Problem data: + problem_vio: wp.array[int32], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + problem_diag: wp.array[float32], + problem_P: wp.array[float32], + problem_v_f: wp.array[float32], + eta: wp.array[float32], + body_space: wp.array[float32], + contact_block_inv: wp.array[mat33f], + block_iteration: int32, + contact_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solution_lambdas: wp.array[float32], +): + _apply_contact_offset_update( + wp.tid(), + num_nzb, + nzb_start, + nzb_coords, + nzb_values, + row_start, + col_start, + contacts_model_active, + contacts_wid, + contacts_cid, + contacts_nzb_offsets, + problem_vio, + problem_nc, + problem_ccgo, + problem_cio, + problem_mu, + problem_diag, + problem_P, + problem_v_f, + eta, + body_space, + contact_block_inv, + block_iteration, + contact_iteration, + solver_config, + solution_lambdas, + ) + + +@wp.kernel +def _solve_dvi_sparse_unilateral_offset_update( + limits_capacity: int32, + # Shared matrix data: + num_nzb: wp.array[int32], + nzb_start: wp.array[int32], + nzb_coords: wp.array2d[int32], + nzb_values: wp.array[vec6f], + row_start: wp.array[int32], + col_start: wp.array[int32], + # Active limits: + limits_model_active: wp.array[int32], + limits_wid: wp.array[int32], + limits_lid: wp.array[int32], + limits_nzb_offsets: wp.array[int32], + problem_nl: wp.array[int32], + problem_lcgo: wp.array[int32], + # Active contacts: + contacts_model_active: wp.array[int32], + contacts_wid: wp.array[int32], + contacts_cid: wp.array[int32], + contacts_nzb_offsets: wp.array[int32], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + contact_block_inv: wp.array[mat33f], + # Shared problem data: + problem_vio: wp.array[int32], + problem_diag: wp.array[float32], + problem_P: wp.array[float32], + problem_v_f: wp.array[float32], + eta: wp.array[float32], + body_space: wp.array[float32], + block_iteration: int32, + contact_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solution_lambdas: wp.array[float32], +): + """Fused limit + contact projection sweep. + + Threads ``[0, limits_capacity)`` update joint limits; the remainder update + contacts. Both read the same ``body_space`` and write disjoint regions of + ``solution_lambdas``, so merging them into one launch is equivalent to the + two separate sweeps while removing a per-iteration kernel launch. + """ + tid = wp.tid() + if tid < limits_capacity: + _apply_limit_offset_update( + tid, + num_nzb, + nzb_start, + nzb_coords, + nzb_values, + row_start, + col_start, + limits_model_active, + limits_wid, + limits_lid, + limits_nzb_offsets, + problem_vio, + problem_nl, + problem_lcgo, + problem_diag, + problem_P, + problem_v_f, + eta, + body_space, + block_iteration, + contact_iteration, + solver_config, + solution_lambdas, + ) + else: + _apply_contact_offset_update( + tid - limits_capacity, + num_nzb, + nzb_start, + nzb_coords, + nzb_values, + row_start, + col_start, + contacts_model_active, + contacts_wid, + contacts_cid, + contacts_nzb_offsets, + problem_vio, + problem_nc, + problem_ccgo, + problem_cio, + problem_mu, + problem_diag, + problem_P, + problem_v_f, + eta, + body_space, + contact_block_inv, + block_iteration, + contact_iteration, + solver_config, + solution_lambdas, + ) + + +@wp.kernel +def _build_sparse_bilateral_block( + # Inputs: + model_bodies_inv_m_i: wp.array[float32], + data_bodies_inv_I_i: wp.array[mat33f], + pair_wid: wp.array[int32], + pair_row: wp.array[int32], + pair_col: wp.array[int32], + pair_bid: wp.array[int32], + pair_i: wp.array[int32], + pair_j: wp.array[int32], + jacobian_cts_nzb_values: wp.array[vec6f], + problem_njc: wp.array[int32], + bilateral_mio: wp.array[int32], + bilateral_vio: wp.array[int32], + bilateral_P: wp.array[float32], + # Output: + bilateral_D: wp.array[float32], +): + pair_id = wp.tid() + wid = pair_wid[pair_id] + njc = problem_njc[wid] + row = pair_row[pair_id] + col = pair_col[pair_id] + block_i = jacobian_cts_nzb_values[pair_i[pair_id]] + block_j = jacobian_cts_nzb_values[pair_j[pair_id]] + Jv_i = vec3f(block_i[0], block_i[1], block_i[2]) + Jv_j = vec3f(block_j[0], block_j[1], block_j[2]) + Jw_i = vec3f(block_i[3], block_i[4], block_i[5]) + Jw_j = vec3f(block_j[3], block_j[4], block_j[5]) + + bid_k = pair_bid[pair_id] + inv_m_k = model_bodies_inv_m_i[bid_k] + inv_I_k = data_bodies_inv_I_i[bid_k] + D_ij = inv_m_k * wp.dot(Jv_i, Jv_j) + wp.dot(Jw_i, inv_I_k @ Jw_j) + + bvio = bilateral_vio[wid] + p_row = bilateral_P[bvio + row] + p_col = bilateral_P[bvio + col] + val = p_row * D_ij * p_col + + bmio = bilateral_mio[wid] + wp.atomic_add(bilateral_D, bmio + njc * row + col, val) + wp.atomic_add(bilateral_D, bmio + njc * col + row, val) + + +@wp.kernel +def _set_sparse_bilateral_diagonal( + # Inputs: + problem_njc: wp.array[int32], + problem_vio: wp.array[int32], + bilateral_mio: wp.array[int32], + bilateral_vio: wp.array[int32], + problem_diag: wp.array[float32], + # Outputs: + bilateral_D: wp.array[float32], + bilateral_P: wp.array[float32], +): + wid, row = wp.tid() + + njc = problem_njc[wid] + if row >= njc: + return + + pvio = problem_vio[wid] + bvio = bilateral_vio[wid] + bmio = bilateral_mio[wid] + diag = wp.abs(problem_diag[pvio + row]) + p = wp.sqrt(1.0 / (diag + FLOAT32_EPS)) + bilateral_P[bvio + row] = p + bilateral_D[bmio + njc * row + row] = p * diag * p + float32(7.0e-7) + + +@wp.kernel +def _compute_sparse_contact_block_inverse( + # Inputs: + model_info_bodies_offset: wp.array[int32], + model_bodies_inv_m_i: wp.array[float32], + data_bodies_inv_I_i: wp.array[mat33f], + jacobian_cts_nzb_start: wp.array[int32], + jacobian_cts_num_nzb: wp.array[int32], + jacobian_cts_nzb_coords: wp.array2d[int32], + jacobian_cts_nzb_values: wp.array[vec6f], + problem_nc: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_vio: wp.array[int32], + problem_P: wp.array[float32], + solver_config: wp.array[DVIConfigStruct], + max_num_nzb: int32, + # Outputs: + contact_block_inv: wp.array[mat33f], +): + wid, cid = wp.tid() + + nc = problem_nc[wid] + cio = problem_cio[wid] + if cid >= nc: + return + + cfg = solver_config[wid] + if not cfg.contact_block_preconditioner: + contact_block_inv[cio + cid] = mat33f(0.0) + return + + ccgo = problem_ccgo[wid] + ccio = ccgo + int32(3) * cid + nzb_start = jacobian_cts_nzb_start[wid] + num_nzb = jacobian_cts_num_nzb[wid] + pvio = problem_vio[wid] + + D = mat33f(0.0) + + for block_id_i in range(max_num_nzb): + if block_id_i >= num_nzb: + continue + + global_block_id_i = nzb_start + block_id_i + block_coords_i = jacobian_cts_nzb_coords[global_block_id_i] + local_i = block_coords_i[0] - ccio + if local_i < 0 or local_i >= int32(3): + continue + + block_i = jacobian_cts_nzb_values[global_block_id_i] + Jv_i = vec3f(block_i[0], block_i[1], block_i[2]) + Jw_i = vec3f(block_i[3], block_i[4], block_i[5]) + p_i = problem_P[pvio + ccio + local_i] + + for block_id_j in range(max_num_nzb): + if block_id_j >= num_nzb: + continue + + global_block_id_j = nzb_start + block_id_j + block_coords_j = jacobian_cts_nzb_coords[global_block_id_j] + local_j = block_coords_j[0] - ccio + if local_j < 0 or local_j >= int32(3) or block_coords_i[1] != block_coords_j[1]: + continue + + block_j = jacobian_cts_nzb_values[global_block_id_j] + Jv_j = vec3f(block_j[0], block_j[1], block_j[2]) + Jw_j = vec3f(block_j[3], block_j[4], block_j[5]) + p_j = problem_P[pvio + ccio + local_j] + + bid = model_info_bodies_offset[wid] + block_coords_i[1] // int32(6) + inv_m = model_bodies_inv_m_i[bid] + inv_I = data_bodies_inv_I_i[bid] + D[local_i, local_j] += p_i * (inv_m * wp.dot(Jv_i, Jv_j) + wp.dot(Jw_i, inv_I @ Jw_j)) * p_j + + D[0, 0] += cfg.regularization + D[1, 1] += cfg.regularization + D[2, 2] += cfg.regularization + + diag_max = wp.max(wp.max(wp.abs(D[0, 0]), wp.abs(D[1, 1])), wp.abs(D[2, 2])) + if diag_max > FLOAT32_EPS: + det = wp.determinant(D) + det_min = FLOAT32_EPS * diag_max * diag_max * diag_max + if det > det_min: + contact_block_inv[cio + cid] = wp.inverse(D) + return + + contact_block_inv[cio + cid] = mat33f(0.0) + + +@wp.kernel +def _solve_dvi_sparse_jacobi_update( + # Inputs: + problem_dim: wp.array[int32], + problem_vio: wp.array[int32], + problem_njc: wp.array[int32], + problem_nl: wp.array[int32], + problem_nc: wp.array[int32], + problem_lcgo: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + problem_diag: wp.array[float32], + problem_P: wp.array[float32], + problem_v_f: wp.array[float32], + state_v_aug: wp.array[float32], + contact_block_inv: wp.array[mat33f], + iteration: int32, + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solution_lambdas: wp.array[float32], +): + wid, tid = wp.tid() + + ncts = problem_dim[wid] + if tid >= ncts: + return + + cfg = solver_config[wid] + if iteration >= cfg.max_iterations: + return + + # Every row uses the same matrix-free D * lambda snapshot. Only the first + # row of each contact updates its three-row block, preserving Jacobi semantics. + vio = problem_vio[wid] + v_i = state_v_aug[vio + tid] + problem_v_f[vio + tid] + P_i = problem_P[vio + tid] + D_ii_raw = wp.abs(problem_diag[vio + tid]) * P_i * P_i + D_ii = D_ii_raw + cfg.regularization + FLOAT32_EPS + + njc = problem_njc[wid] + if tid < njc: + solution_lambdas[vio + tid] += -cfg.omega * v_i / D_ii + return + + nl = problem_nl[wid] + lcgo = problem_lcgo[wid] + lid = tid - lcgo + if lid >= 0 and lid < nl: + lambda_limit_old = solution_lambdas[vio + tid] + lambda_limit_new = lambda_limit_old + if D_ii_raw > FLOAT32_EPS: + lambda_limit_new = wp.max(0.0, lambda_limit_old - cfg.omega * v_i / D_ii) + solution_lambdas[vio + tid] = lambda_limit_new + return + + nc = problem_nc[wid] + ccgo = problem_ccgo[wid] + contact_row = tid - ccgo + if contact_row < 0 or contact_row >= int32(3) * nc or contact_row % int32(3) != int32(0): + return + + cid = contact_row // int32(3) + ccio = ccgo + int32(3) * cid + ccio_v = vio + ccio + mu_c = problem_mu[problem_cio[wid] + cid] + + v_t0 = state_v_aug[ccio_v + 0] + problem_v_f[ccio_v + 0] + v_t1 = state_v_aug[ccio_v + 1] + problem_v_f[ccio_v + 1] + v_n = state_v_aug[ccio_v + 2] + problem_v_f[ccio_v + 2] + mu_c * wp.sqrt(v_t0 * v_t0 + v_t1 * v_t1) + v_c = vec3f(v_t0, v_t1, v_n) + + P_0 = problem_P[ccio_v + 0] + P_1 = problem_P[ccio_v + 1] + P_2 = problem_P[ccio_v + 2] + D_00 = wp.abs(problem_diag[ccio_v + 0]) * P_0 * P_0 + D_11 = wp.abs(problem_diag[ccio_v + 1]) * P_1 * P_1 + D_22 = wp.abs(problem_diag[ccio_v + 2]) * P_2 * P_2 + + lambda_contact_old = vec3f( + solution_lambdas[ccio_v + 0], + solution_lambdas[ccio_v + 1], + solution_lambdas[ccio_v + 2], + ) + D_diag = _contact_trace_preconditioner(vec3f(D_00, D_11, D_22)) + if cfg.contact_block_preconditioner: + lambda_projected = _project_contact_block_update( + lambda_contact_old, + v_c, + D_diag, + contact_block_inv[problem_cio[wid] + cid], + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + else: + lambda_projected = _project_contact_diagonal_update( + lambda_contact_old, + v_c, + D_diag, + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + lambda_contact_new = lambda_contact_old + cfg.contact_jacobi_relaxation * (lambda_projected - lambda_contact_old) + + solution_lambdas[ccio_v + 0] = lambda_contact_new.x + solution_lambdas[ccio_v + 1] = lambda_contact_new.y + solution_lambdas[ccio_v + 2] = lambda_contact_new.z + + +@wp.kernel +def _solve_dvi_sparse_unilateral_jacobi_update( + # Inputs: + problem_dim: wp.array[int32], + problem_vio: wp.array[int32], + problem_njc: wp.array[int32], + problem_nl: wp.array[int32], + problem_nc: wp.array[int32], + problem_lcgo: wp.array[int32], + problem_ccgo: wp.array[int32], + problem_cio: wp.array[int32], + problem_mu: wp.array[float32], + problem_diag: wp.array[float32], + problem_P: wp.array[float32], + problem_v_f: wp.array[float32], + state_v_aug: wp.array[float32], + contact_block_inv: wp.array[mat33f], + block_iteration: int32, + contact_iteration: int32, + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solution_lambdas: wp.array[float32], +): + wid, tid = wp.tid() + + ncts = problem_dim[wid] + if tid >= ncts: + return + + cfg = solver_config[wid] + if block_iteration >= cfg.block_iterations or contact_iteration >= cfg.contact_iterations: + return + + njc = problem_njc[wid] + if tid < njc: + return + + vio = problem_vio[wid] + v_i = state_v_aug[vio + tid] + problem_v_f[vio + tid] + P_i = problem_P[vio + tid] + D_ii_raw = wp.abs(problem_diag[vio + tid]) * P_i * P_i + D_ii = D_ii_raw + cfg.regularization + FLOAT32_EPS + + nl = problem_nl[wid] + lcgo = problem_lcgo[wid] + lid = tid - lcgo + if lid >= 0 and lid < nl: + lambda_limit_old = solution_lambdas[vio + tid] + lambda_limit_new = lambda_limit_old + if D_ii_raw > FLOAT32_EPS: + lambda_limit_new = wp.max(0.0, lambda_limit_old - cfg.omega * v_i / D_ii) + solution_lambdas[vio + tid] = lambda_limit_new + return + + nc = problem_nc[wid] + ccgo = problem_ccgo[wid] + contact_row = tid - ccgo + if contact_row < 0 or contact_row >= int32(3) * nc or contact_row % int32(3) != int32(0): + return + + cid = contact_row // int32(3) + ccio = ccgo + int32(3) * cid + ccio_v = vio + ccio + mu_c = problem_mu[problem_cio[wid] + cid] + + v_t0 = state_v_aug[ccio_v + 0] + problem_v_f[ccio_v + 0] + v_t1 = state_v_aug[ccio_v + 1] + problem_v_f[ccio_v + 1] + v_n = state_v_aug[ccio_v + 2] + problem_v_f[ccio_v + 2] + mu_c * wp.sqrt(v_t0 * v_t0 + v_t1 * v_t1) + v_c = vec3f(v_t0, v_t1, v_n) + + P_0 = problem_P[ccio_v + 0] + P_1 = problem_P[ccio_v + 1] + P_2 = problem_P[ccio_v + 2] + D_00 = wp.abs(problem_diag[ccio_v + 0]) * P_0 * P_0 + D_11 = wp.abs(problem_diag[ccio_v + 1]) * P_1 * P_1 + D_22 = wp.abs(problem_diag[ccio_v + 2]) * P_2 * P_2 + + lambda_contact_old = vec3f( + solution_lambdas[ccio_v + 0], + solution_lambdas[ccio_v + 1], + solution_lambdas[ccio_v + 2], + ) + D_diag = _contact_trace_preconditioner(vec3f(D_00, D_11, D_22)) + if cfg.contact_block_preconditioner: + lambda_projected = _project_contact_block_update( + lambda_contact_old, + v_c, + D_diag, + contact_block_inv[problem_cio[wid] + cid], + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + else: + lambda_projected = _project_contact_diagonal_update( + lambda_contact_old, + v_c, + D_diag, + cfg.regularization, + cfg.contact_jacobi_omega, + mu_c, + ) + lambda_contact_new = lambda_contact_old + cfg.contact_jacobi_relaxation * (lambda_projected - lambda_contact_old) + + solution_lambdas[ccio_v + 0] = lambda_contact_new.x + solution_lambdas[ccio_v + 1] = lambda_contact_new.y + solution_lambdas[ccio_v + 2] = lambda_contact_new.z + + +@wp.kernel +def _compute_dvi_sparse_solution_vectors( + # Inputs: + problem_dim: wp.array[int32], + problem_vio: wp.array[int32], + problem_v_f: wp.array[float32], + # Outputs: + state_s: wp.array[float32], + state_v_aug: wp.array[float32], + solution_v_plus: wp.array[float32], +): + wid, tid = wp.tid() + + ncts = problem_dim[wid] + if tid >= ncts: + return + + v_i = problem_vio[wid] + tid + v_plus = state_v_aug[v_i] + problem_v_f[v_i] + solution_v_plus[v_i] = v_plus + state_v_aug[v_i] = v_plus + state_s[v_i] = 0.0 + + +@wp.kernel +def _set_dvi_sparse_status_iterations( + # Inputs: + problem_dim: wp.array[int32], + solver_config: wp.array[DVIConfigStruct], + # Outputs: + solver_status: wp.array[DVIStatus], +): + wid = wp.tid() + status = solver_status[wid] + if problem_dim[wid] == int32(0): + status.iterations = int32(0) + else: + # Sparse Jacobi currently runs a fixed iteration count; terminal DVI + # residuals are computed by the shared dense/sparse status kernel. + status.iterations = solver_config[wid].max_iterations + solver_status[wid] = status diff --git a/newton/_src/solvers/kamino/_src/solvers/dvi/types.py b/newton/_src/solvers/kamino/_src/solvers/dvi/types.py new file mode 100644 index 0000000000..0714dbee48 --- /dev/null +++ b/newton/_src/solvers/kamino/_src/solvers/dvi/types.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Data containers for the Kamino DVI solver.""" + +from __future__ import annotations + +import warp as wp + +from ....config import DVISolverConfig +from ...core.size import SizeKamino +from ...linalg import DenseLinearOperatorData +from ..common import DualSolution + +wp.set_module_options({"enable_backward": False}) + +float32 = wp.float32 +int32 = wp.int32 +mat33f = wp.mat33f +vec2f = wp.vec2f + + +@wp.struct +class DVIConfigStruct: + """On-device DVI solver configuration.""" + + tolerance: float32 + """Tolerance for iterate-change stopping and terminal DVI residuals.""" + + regularization: float32 + """Diagonal regularization used by projected Gauss-Seidel updates.""" + + omega: float32 + """Projected Gauss-Seidel update relaxation.""" + + max_iterations: int32 + """Maximum projected Gauss-Seidel iterations for the fallback path.""" + + block_iterations: int32 + """Outer direct-bilateral/projected-inequality block iterations.""" + + contact_iterations: int32 + """Projected sweeps for unilateral inequalities in each direct-bilateral block.""" + + bilateral_solve_period: int32 + """Block iteration period for repeated direct bilateral solves.""" + + contact_jacobi_omega: float32 + """Step size for contact Jacobi and block-preconditioned updates.""" + + contact_jacobi_relaxation: float32 + """Solution mixing for contact Jacobi and block-preconditioned updates.""" + + contact_block_preconditioner: wp.bool + """Whether to use the full contact 3x3 diagonal block for projected updates.""" + + +@wp.struct +class DVIStatus: + """Per-world DVI convergence status.""" + + converged: int32 + """Whether all terminal feasibility, equality, and complementarity residuals satisfy tolerance.""" + iterations: int32 + """Projected sweeps; direct-bilateral solves report block/contact sweeps.""" + r_p: float32 + """Primal cone-feasibility residual.""" + r_d: float32 + """Maximum dual cone-feasibility and bilateral velocity residual.""" + r_c: float32 + """Maximum absolute impulse-velocity inner product.""" + r_b: float32 + """Bilateral constraint-space velocity residual.""" + + +class DVIInfo: + """Optional terminal convergence diagnostics for each simulated world.""" + + def __init__(self, size: SizeKamino | None = None): + self.status: wp.array[DVIStatus] | None = None + """Terminal DVI status, shape ``(num_worlds,)``.""" + if size is not None: + self.finalize(size) + + def finalize(self, size: SizeKamino) -> None: + """Allocate diagnostic arrays for a model size.""" + self.status = wp.zeros(shape=(size.num_worlds,), dtype=DVIStatus) + + def zero(self) -> None: + """Reset diagnostics to zero.""" + self.status.zero_() + + +class DVIState: + """Scratch arrays used by the DVI solver.""" + + def __init__(self, size: SizeKamino | None = None): + self.sigma: wp.array[vec2f] | None = None + """Zero proximal terms used when evaluating shared solution metrics.""" + self.v_aug: wp.array[float32] | None = None + self.s: wp.array[float32] | None = None + self.scratch: wp.array[float32] | None = None + self.bilateral_rhs: wp.array[float32] | None = None + self.bilateral_solution: wp.array[float32] | None = None + self.bilateral_preconditioner: wp.array[float32] | None = None + self.bilateral_active_dim: wp.array[int32] | None = None + self.contact_block_inv: wp.array[mat33f] | None = None + self.contact_colors: wp.array[int32] | None = None + self.contact_num_colors: wp.array[int32] | None = None + if size is not None: + self.finalize(size) + + def finalize(self, size: SizeKamino): + """Allocate scratch arrays for the supplied model size.""" + self.sigma = wp.zeros(size.num_worlds, dtype=vec2f) + self.v_aug = wp.zeros(size.sum_of_max_total_cts, dtype=float32) + self.s = wp.zeros(size.sum_of_max_total_cts, dtype=float32) + self.scratch = wp.zeros(size.sum_of_max_total_cts, dtype=float32) + self.bilateral_rhs = wp.zeros(size.sum_of_num_joint_cts, dtype=float32) + self.bilateral_solution = wp.zeros(size.sum_of_num_joint_cts, dtype=float32) + self.bilateral_preconditioner = wp.zeros(size.sum_of_num_joint_cts, dtype=float32) + self.bilateral_active_dim = wp.zeros(size.num_worlds, dtype=int32) + self.contact_block_inv = wp.zeros(max(1, size.sum_of_max_contacts), dtype=mat33f) + self.contact_colors = wp.full(max(1, size.sum_of_max_contacts), -1, dtype=int32) + self.contact_num_colors = wp.zeros(max(1, size.num_worlds), dtype=int32) + + def reset(self): + """Reset scratch arrays to zero.""" + self.sigma.zero_() + self.v_aug.zero_() + self.s.zero_() + self.scratch.zero_() + self.bilateral_rhs.zero_() + self.bilateral_solution.zero_() + self.bilateral_preconditioner.zero_() + self.bilateral_active_dim.zero_() + self.contact_block_inv.zero_() + self.contact_colors.fill_(-1) + self.contact_num_colors.zero_() + + +class DVIData: + """High-level DVI solver data.""" + + def __init__( + self, + size: SizeKamino | None = None, + collect_info: bool = False, + device: wp.DeviceLike = None, + ): + self.config: wp.array[DVIConfigStruct] | None = None + self.status: wp.array[DVIStatus] | None = None + self.state: DVIState | None = None + self.solution: DualSolution | None = None + self.info: DVIInfo | None = None + self.bilateral_operator: DenseLinearOperatorData | None = None + if size is not None: + self.finalize(size=size, collect_info=collect_info, device=device) + + def finalize(self, size: SizeKamino, collect_info: bool = False, device: wp.DeviceLike = None): + """Allocate DVI data arrays.""" + with wp.ScopedDevice(device): + self.config = wp.zeros(shape=(size.num_worlds,), dtype=DVIConfigStruct) + self.status = wp.zeros(shape=(size.num_worlds,), dtype=DVIStatus) + self.state = DVIState(size) + self.solution = DualSolution(size) + self.info = DVIInfo(size) if collect_info else None + self.bilateral_operator = None + + +def convert_config_to_struct(config: DVISolverConfig) -> DVIConfigStruct: + """Convert a host-side DVI config to an on-device struct.""" + config_struct = DVIConfigStruct() + config_struct.tolerance = config.tolerance + config_struct.regularization = config.regularization + config_struct.omega = config.omega + config_struct.max_iterations = config.max_iterations + config_struct.block_iterations = config.block_iterations + config_struct.contact_iterations = config.contact_iterations + config_struct.bilateral_solve_period = config.bilateral_solve_period + config_struct.contact_jacobi_omega = config.contact_jacobi_omega + config_struct.contact_jacobi_relaxation = config.contact_jacobi_relaxation + config_struct.contact_block_preconditioner = config.contact_block_preconditioner + return config_struct diff --git a/newton/_src/solvers/kamino/_src/solvers/fk/kernels.py b/newton/_src/solvers/kamino/_src/solvers/fk/kernels.py index c9b9d10a3f..15a956f243 100644 --- a/newton/_src/solvers/kamino/_src/solvers/fk/kernels.py +++ b/newton/_src/solvers/kamino/_src/solvers/fk/kernels.py @@ -11,18 +11,21 @@ from ...core.joints import JointActuationType, JointDoFType from ...core.math import ( - TWO_PI, G_of, quat_left_jacobian_inverse, quat_log, - squared_norm, unit_quat_apply, unit_quat_apply_jacobian, unit_quat_conj_apply, unit_quat_conj_apply_jacobian, unit_quat_conj_to_rotation_matrix, ) -from ...kinematics.joints import get_joint_coords_mapping_function +from ...core.types import mat34f +from ...kinematics.joints import ( + correct_quat_vector_coord, + correct_rotational_coord, + get_joint_coords_mapping_function, +) from ...linalg.sparse_matrix import BlockDType from .types import FKJointDoFType @@ -33,6 +36,8 @@ __all__ = [ "_add_regularizer_to_diagonal", "_apply_line_search_step", + "_compute_fk_axis_joint_frames", + "_compute_fk_joint_frames", "_correct_actuator_coords", "_correct_universal_constraint_velocities", "_eval_actuator_coords", @@ -53,6 +58,7 @@ "_newton_check", "_reset_state", "_reset_state_base_q", + "_resolve_fk_actuation_types", "_update_cg_tolerance_kernel", "create_1d_tile_based_kernels", "create_2d_tile_based_kernels", @@ -61,6 +67,7 @@ "create_eval_joint_constraints_sparse_jacobian_kernel", "create_eval_min_num_iterations_kernel", "read_quat_from_array", + "validate_fk_actuation_updates", ] @@ -96,11 +103,268 @@ def read_quat_from_array(array: wp.array[wp.float32], offset: int, normalize: bo return q +@wp.func +def _resolve_fk_actuation_type(act_type: wp.int32, fk_act_flag: wp.int32) -> wp.int32: + """Apply an optional FK actuation override to a model actuation type.""" + if fk_act_flag == 0: + return JointActuationType.PASSIVE + if fk_act_flag == 1: + return JointActuationType.FORCE + return act_type + + +@wp.func +def _load_joint_poses( + base_id: wp.int32, follower_id: wp.int32, bodies_q: wp.array[wp.transformf] +) -> tuple[wp.vec3f, wp.quatf, wp.vec3f, wp.quatf]: + """Load the base and follower poses, using the identity pose for the world.""" + c_base = wp.vec3f(0.0, 0.0, 0.0) + q_base = wp.quatf(0.0, 0.0, 0.0, 1.0) + if base_id >= 0: + c_base = wp.transform_get_translation(bodies_q[base_id]) + q_base = wp.transform_get_rotation(bodies_q[base_id]) + c_follower = wp.transform_get_translation(bodies_q[follower_id]) + q_follower = wp.transform_get_rotation(bodies_q[follower_id]) + return c_base, q_base, c_follower, q_follower + + +@wp.func +def _get_reduced_constraint_ids( + joint_id: wp.int32, ct_full_to_red_map: wp.array[wp.int32] +) -> tuple[wp.vec3i, wp.vec3i]: + """Load the translational and rotational reduced constraint ids for a joint.""" + first_ct_id_full = 6 * joint_id + trans_ct_ids_red = wp.vec3i( + ct_full_to_red_map[first_ct_id_full], + ct_full_to_red_map[first_ct_id_full + 1], + ct_full_to_red_map[first_ct_id_full + 2], + ) + rot_ct_ids_red = wp.vec3i( + ct_full_to_red_map[first_ct_id_full + 3], + ct_full_to_red_map[first_ct_id_full + 4], + ct_full_to_red_map[first_ct_id_full + 5], + ) + return trans_ct_ids_red, rot_ct_ids_red + + +@wp.func +def _eval_translation_jacobian_blocks( + X_T: wp.mat33f, + q_base: wp.quatf, + q_follower: wp.quatf, + x_follower: wp.vec3f, + c_base: wp.vec3f, + c_follower: wp.vec3f, + has_base: wp.bool, +) -> tuple[wp.mat33f, mat34f, wp.mat33f, mat34f]: + """Evaluate the base and follower blocks of a translational constraint Jacobian.""" + X_T_R_base_T = X_T * unit_quat_conj_to_rotation_matrix(q_base) + jac_trans_c_base = wp.mat33f(0.0) + jac_trans_q_base = mat34f(0.0) + if has_base: + jac_trans_c_base = -X_T_R_base_T + delta_pos = unit_quat_apply(q_follower, x_follower) + c_follower - c_base + jac_trans_q_base = X_T * unit_quat_conj_apply_jacobian(q_base, delta_pos) + jac_trans_c_follower = X_T_R_base_T + jac_trans_q_follower = X_T_R_base_T * unit_quat_apply_jacobian(q_follower, x_follower) + return jac_trans_c_base, jac_trans_q_base, jac_trans_c_follower, jac_trans_q_follower + + +@wp.func +def _eval_rotation_jacobian_blocks( + X_T: wp.mat33f, + q_base: wp.quatf, + q_follower: wp.quatf, + q_rel_body: wp.quatf, + has_base: wp.bool, +) -> tuple[mat34f, mat34f]: + """Evaluate the base and follower blocks of a rotational constraint Jacobian.""" + q_base_sq_norm = wp.dot(q_base, q_base) + q_follower_sq_norm = wp.dot(q_follower, q_follower) + R_base_T = unit_quat_conj_to_rotation_matrix(q_base / wp.sqrt(q_base_sq_norm)) + q_rel = q_follower * wp.quat_inverse(q_rel_body) * wp.quat_inverse(q_base) + temp = X_T * R_base_T * quat_left_jacobian_inverse(q_rel) + jac_rot_q_base = mat34f(0.0) + if has_base: + jac_rot_q_base = (-2.0 / q_base_sq_norm) * temp * G_of(q_base) + jac_rot_q_follower = (2.0 / q_follower_sq_norm) * temp * G_of(q_follower) + return jac_rot_q_base, jac_rot_q_follower + + +@wp.func +def _eval_passive_universal_jacobian_blocks( + X_T: wp.mat33f, q_base: wp.quatf, q_follower: wp.quatf, has_base: wp.bool +) -> tuple[wp.vec4f, wp.vec4f]: + """Evaluate the base and follower blocks of a passive universal constraint Jacobian.""" + a_x = X_T[0] + a_y = X_T[1] + jac_q_base = wp.vec4f(0.0) + if has_base: + a_y_follower = unit_quat_apply(q_follower, a_y) + jac_q_base = -a_y_follower * unit_quat_apply_jacobian(q_base, a_x) + a_x_base = unit_quat_apply(q_base, a_x) + jac_q_follower = -a_x_base * unit_quat_apply_jacobian(q_follower, a_y) + return jac_q_base, jac_q_follower + + +@wp.func +def _correct_rotational_actuator_coord( + actuators_q: wp.array[wp.float32], actuators_q_ref: wp.array[wp.float32], coord_id: wp.int32 +): + """Correct an angular actuator coordinate against its reference.""" + actuators_q[coord_id] = correct_rotational_coord(actuators_q[coord_id], actuators_q_ref[coord_id]) + + +@wp.func +def _correct_quat_actuator_coords( + actuators_q: wp.array[wp.float32], actuators_q_ref: wp.array[wp.float32], coord_id: wp.int32 +): + """Correct four quaternion actuator coordinates against their reference.""" + quat = wp.vec4f( + actuators_q[coord_id], actuators_q[coord_id + 1], actuators_q[coord_id + 2], actuators_q[coord_id + 3] + ) + quat_ref = wp.vec4f( + actuators_q_ref[coord_id], + actuators_q_ref[coord_id + 1], + actuators_q_ref[coord_id + 2], + actuators_q_ref[coord_id + 3], + ) + quat_corrected = correct_quat_vector_coord(quat, quat_ref) + for i in range(4): + actuators_q[coord_id + i] = quat_corrected[i] + + ### # Kernels ### +@wp.kernel +def _resolve_fk_actuation_types( + # Inputs + model_act_type: wp.array[wp.int32], + model_fk_act_flag: wp.array[wp.int32], + # Outputs + fk_act_type: wp.array[wp.int32], +): + """Resolve effective FK actuation types for the main model joints.""" + joint = wp.tid() + flag = wp.int32(-1) + if model_fk_act_flag: + flag = model_fk_act_flag[joint] + fk_act_type[joint] = _resolve_fk_actuation_type(model_act_type[joint], flag) + + +@wp.kernel +def validate_fk_actuation_updates( + # Inputs + model_act_type: wp.array[wp.int32], + model_fk_act_flag: wp.array[wp.int32], + built_fk_actuated: wp.array[wp.int32], + # Outputs + violations: wp.array[wp.int32], +): + """Find invalid overrides and changes to the set of joints actuated for FK. + + ``built_fk_actuated`` is indexed by model joint: 0 is passive, 1 is + actuated, and -1 skips validation for a joint replaced by FK. + """ + joint = wp.tid() + flag = wp.int32(-1) + if model_fk_act_flag: + flag = model_fk_act_flag[joint] + if flag < -1 or flag > 1: + wp.atomic_min(violations, 1, joint) + return + + built_actuated = built_fk_actuated[joint] + if built_actuated < 0: + return + + act_type = _resolve_fk_actuation_type(model_act_type[joint], flag) + if (act_type != JointActuationType.PASSIVE) != (built_actuated != 0): + wp.atomic_min(violations, 0, joint) + + +@wp.kernel +def _compute_fk_joint_frames( + # Inputs + source_joint: wp.array[wp.int32], + model_B_r_Bj: wp.array[wp.vec3f], + model_F_r_Fj: wp.array[wp.vec3f], + model_X_Bj: wp.array[wp.mat33f], + model_X_Fj: wp.array[wp.mat33f], + # Outputs + fk_B_r_Bj: wp.array[wp.vec3f], + fk_F_r_Fj: wp.array[wp.vec3f], + fk_X_Bj: wp.array[wp.mat33f], + fk_X_Fj: wp.array[wp.mat33f], +): + """Compute FK joint frames from the current model data.""" + fk_joint = wp.tid() + model_joint = source_joint[fk_joint] + if model_joint >= 0: + # Preserve frames for joints copied from the model. + fk_B_r_Bj[fk_joint] = model_B_r_Bj[model_joint] + fk_F_r_Fj[fk_joint] = model_F_r_Fj[model_joint] + fk_X_Bj[fk_joint] = model_X_Bj[model_joint] + fk_X_Fj[fk_joint] = model_X_Fj[model_joint] + else: + # FK-added base and axis joints start at identity; axis orientations are + # overwritten by _compute_fk_axis_joint_frames. + fk_B_r_Bj[fk_joint] = wp.vec3f(0.0) + fk_F_r_Fj[fk_joint] = wp.vec3f(0.0) + fk_X_Bj[fk_joint] = wp.identity(n=3, dtype=wp.float32) + fk_X_Fj[fk_joint] = wp.identity(n=3, dtype=wp.float32) + + +@wp.kernel +def _compute_fk_axis_joint_frames( + # Inputs + axis_fk_joint: wp.array[wp.int32], + axis_body: wp.array[wp.int32], + axis_joint_0: wp.array[wp.int32], + axis_joint_1: wp.array[wp.int32], + model_joint_bid_B: wp.array[wp.int32], + model_joint_B_r_Bj: wp.array[wp.vec3f], + model_joint_F_r_Fj: wp.array[wp.vec3f], + model_body_q_0: wp.array[wp.transformf], + # Outputs + fk_X_Bj: wp.array[wp.mat33f], + fk_X_Fj: wp.array[wp.mat33f], +): + """Compute synthetic axis-joint frames from the model data.""" + axis_joint = wp.tid() + fk_joint = axis_fk_joint[axis_joint] + body = axis_body[axis_joint] + joint_0 = axis_joint_0[axis_joint] + joint_1 = axis_joint_1[axis_joint] + body_q = model_body_q_0[body] + + # Locate both spherical-joint anchors in the tie-rod body frame. + local_0 = model_joint_F_r_Fj[joint_0] + if model_joint_bid_B[joint_0] == body: + local_0 = model_joint_B_r_Bj[joint_0] + local_1 = model_joint_F_r_Fj[joint_1] + if model_joint_bid_B[joint_1] == body: + local_1 = model_joint_B_r_Bj[joint_1] + + # Evaluate the anchors in the initial pose and align the joint X axis with + # the line that connects them. + pos_0 = wp.transform_point(body_q, local_0) + pos_1 = wp.transform_point(body_q, local_1) + a_x = wp.normalize(pos_1 - pos_0) + if wp.abs(a_x[2]) < 0.99: + a_y = wp.normalize(wp.cross(wp.vec3f(0.0, 0.0, 1.0), a_x)) + else: + a_y = wp.normalize(wp.cross(wp.vec3f(0.0, 1.0, 0.0), a_x)) + a_z = wp.normalize(wp.cross(a_x, a_y)) + X_Bj = wp.matrix_from_cols(a_x, a_y, a_z) + fk_X_Bj[fk_joint] = X_Bj + # Match the follower frame to the base frame in the initial pose. + fk_X_Fj[fk_joint] = wp.quat_to_matrix(wp.transform_get_rotation(body_q)) * X_Bj + + @wp.kernel def _reset_state( # Inputs @@ -339,15 +603,8 @@ def _eval_actuator_coords( # Get base and follower transformations base_id = joints_bid_B[jt_id] - if base_id < 0: - c_base = wp.vec3f(0.0, 0.0, 0.0) - q_base = wp.quatf(0.0, 0.0, 0.0, 1.0) - else: - c_base = wp.transform_get_translation(bodies_q[base_id]) - q_base = wp.transform_get_rotation(bodies_q[base_id]) follower_id = joints_bid_F[jt_id] - c_follower = wp.transform_get_translation(bodies_q[follower_id]) - q_follower = wp.transform_get_rotation(bodies_q[follower_id]) + c_base, q_base, c_follower, q_follower = _load_joint_poses(base_id, follower_id, bodies_q) # Compute relative pose of follower body in joint frame of base body pos_base = c_base + wp.quat_rotate(q_base, x_base) @@ -361,20 +618,6 @@ def _eval_actuator_coords( _joint_transform_to_coords(dof_type, pos_rel, q_rel, coord_id, actuators_q) -@wp.func -def _correct_joint_angle(angle: wp.float32, angle_ref: wp.float32) -> wp.float32: - """Function adding multiples of 2 pi to an angle, so that it is the closest to a reference.""" - return angle + wp.round((angle_ref - angle) / TWO_PI) * TWO_PI - - -@wp.func -def _correct_joint_quaternion(quat: wp.vec4f, quat_ref: wp.vec4f) -> wp.vec4f: - """Function flipping the sign of a quaternion if needed, so it is the closest to a reference.""" - if squared_norm(quat + quat_ref) < squared_norm(quat - quat_ref): - return -quat - return quat - - @wp.kernel def _correct_actuator_coords( # Inputs @@ -412,46 +655,16 @@ def _correct_actuator_coords( ): # No correction needed return elif dof_type == FKJointDoFType.CYLINDRICAL: # Correct angle up to +/- 2 pi - angle = actuators_q[coord_id + 1] - angle_ref = actuators_q_ref[coord_id + 1] - actuators_q[coord_id + 1] = _correct_joint_angle(angle, angle_ref) + _correct_rotational_actuator_coord(actuators_q, actuators_q_ref, coord_id + 1) elif dof_type == FKJointDoFType.FREE: # Correct quaternion up to sign - quat = wp.vec4f( - actuators_q[coord_id + 3], actuators_q[coord_id + 4], actuators_q[coord_id + 5], actuators_q[coord_id + 6] - ) - quat_ref = wp.vec4f( - actuators_q_ref[coord_id + 3], - actuators_q_ref[coord_id + 4], - actuators_q_ref[coord_id + 5], - actuators_q_ref[coord_id + 6], - ) - quat_corrected = _correct_joint_quaternion(quat, quat_ref) - for i in range(4): - actuators_q[coord_id + 3 + i] = quat_corrected[i] + _correct_quat_actuator_coords(actuators_q, actuators_q_ref, coord_id + 3) elif dof_type == FKJointDoFType.REVOLUTE: # Correct angle up to +/- 2 pi - angle = actuators_q[coord_id] - angle_ref = actuators_q_ref[coord_id] - actuators_q[coord_id] = _correct_joint_angle(angle, angle_ref) + _correct_rotational_actuator_coord(actuators_q, actuators_q_ref, coord_id) elif dof_type == FKJointDoFType.SPHERICAL: # Correct quaternion up to sign - quat = wp.vec4f( - actuators_q[coord_id], actuators_q[coord_id + 1], actuators_q[coord_id + 2], actuators_q[coord_id + 3] - ) - quat_ref = wp.vec4f( - actuators_q_ref[coord_id], - actuators_q_ref[coord_id + 1], - actuators_q_ref[coord_id + 2], - actuators_q_ref[coord_id + 3], - ) - quat_corrected = _correct_joint_quaternion(quat, quat_ref) - for i in range(4): - actuators_q[coord_id + i] = quat_corrected[i] + _correct_quat_actuator_coords(actuators_q, actuators_q_ref, coord_id) elif dof_type == FKJointDoFType.UNIVERSAL: # Correct angles up to +/- 2 pi - angle = actuators_q[coord_id] - angle_ref = actuators_q_ref[coord_id] - actuators_q[coord_id] = _correct_joint_angle(angle, angle_ref) - angle = actuators_q[coord_id + 1] - angle_ref = actuators_q_ref[coord_id + 1] - actuators_q[coord_id + 1] = _correct_joint_angle(angle, angle_ref) + _correct_rotational_actuator_coord(actuators_q, actuators_q_ref, coord_id) + _correct_rotational_actuator_coord(actuators_q, actuators_q_ref, coord_id + 1) else: assert False, "Unexpected actuator dof type" # noqa: B011 @@ -801,17 +1014,7 @@ def _eval_joint_constraints( jt_id_tot = first_joint_id[wd_id] + jt_id_loc # Get reduced constraint ids (-1 meaning constraint is not used) - first_ct_id_full = 6 * jt_id_tot - trans_ct_ids_red = wp.vec3i( - ct_full_to_red_map[first_ct_id_full], - ct_full_to_red_map[first_ct_id_full + 1], - ct_full_to_red_map[first_ct_id_full + 2], - ) - rot_ct_ids_red = wp.vec3i( - ct_full_to_red_map[first_ct_id_full + 3], - ct_full_to_red_map[first_ct_id_full + 4], - ct_full_to_red_map[first_ct_id_full + 5], - ) + trans_ct_ids_red, rot_ct_ids_red = _get_reduced_constraint_ids(jt_id_tot, ct_full_to_red_map) # Get joint local positions and orientation x_base = joints_B_r_B[jt_id_tot] @@ -820,15 +1023,8 @@ def _eval_joint_constraints( # Get base and follower transformations base_id = joints_bid_B[jt_id_tot] - if base_id < 0: - c_base = wp.vec3f(0.0, 0.0, 0.0) - q_base = wp.quatf(0.0, 0.0, 0.0, 1.0) - else: - c_base = wp.transform_get_translation(bodies_q[base_id]) - q_base = wp.transform_get_rotation(bodies_q[base_id]) follower_id = joints_bid_F[jt_id_tot] - c_follower = wp.transform_get_translation(bodies_q[follower_id]) - q_follower = wp.transform_get_rotation(bodies_q[follower_id]) + c_base, q_base, c_follower, q_follower = _load_joint_poses(base_id, follower_id, bodies_q) # Get target relative transformation, in joint/body frame for translation/rotation part t_rel_joint = wp.transform_get_translation(target_rel_transforms[jt_id_tot]) @@ -1013,17 +1209,7 @@ def _eval_joint_constraints_jacobian( jt_id_tot = first_joint_id[wd_id] + jt_id_loc # Get reduced constraint ids (-1 meaning constraint is not used) - first_ct_id_full = 6 * jt_id_tot - trans_ct_ids_red = wp.vec3i( - ct_full_to_red_map[first_ct_id_full], - ct_full_to_red_map[first_ct_id_full + 1], - ct_full_to_red_map[first_ct_id_full + 2], - ) - rot_ct_ids_red = wp.vec3i( - ct_full_to_red_map[first_ct_id_full + 3], - ct_full_to_red_map[first_ct_id_full + 4], - ct_full_to_red_map[first_ct_id_full + 5], - ) + trans_ct_ids_red, rot_ct_ids_red = _get_reduced_constraint_ids(jt_id_tot, ct_full_to_red_map) # Get joint local positions and orientation x_follower = joints_F_r_F[jt_id_tot] @@ -1031,15 +1217,8 @@ def _eval_joint_constraints_jacobian( # Get base and follower transformations base_id_tot = joints_bid_B[jt_id_tot] - if base_id_tot < 0: - c_base = wp.vec3f(0.0, 0.0, 0.0) - q_base = wp.quatf(0.0, 0.0, 0.0, 1.0) - else: - c_base = wp.transform_get_translation(bodies_q[base_id_tot]) - q_base = wp.transform_get_rotation(bodies_q[base_id_tot]) follower_id_tot = joints_bid_F[jt_id_tot] - c_follower = wp.transform_get_translation(bodies_q[follower_id_tot]) - q_follower = wp.transform_get_rotation(bodies_q[follower_id_tot]) + c_base, q_base, c_follower, q_follower = _load_joint_poses(base_id_tot, follower_id_tot, bodies_q) base_id_loc = base_id_tot - first_body_id[wd_id] follower_id_loc = follower_id_tot - first_body_id[wd_id] @@ -1047,23 +1226,15 @@ def _eval_joint_constraints_jacobian( q_rel_body = wp.transform_get_rotation(target_rel_transforms[jt_id_tot]) # Translation constraints - X_T_R_base_T = X_T * unit_quat_conj_to_rotation_matrix(q_base) - if base_id_tot >= 0: - jac_trans_c_base = -X_T_R_base_T - delta_pos = unit_quat_apply(q_follower, x_follower) + c_follower - c_base - jac_trans_q_base = X_T * unit_quat_conj_apply_jacobian(q_base, delta_pos) - jac_trans_c_follower = X_T_R_base_T - jac_trans_q_follower = X_T_R_base_T * unit_quat_apply_jacobian(q_follower, x_follower) - + jac_trans_c_base, jac_trans_q_base, jac_trans_c_follower, jac_trans_q_follower = ( + _eval_translation_jacobian_blocks( + X_T, q_base, q_follower, x_follower, c_base, c_follower, base_id_tot >= 0 + ) + ) # Rotation constraints - q_base_sq_norm = wp.dot(q_base, q_base) - q_follower_sq_norm = wp.dot(q_follower, q_follower) - R_base_T = unit_quat_conj_to_rotation_matrix(q_base / wp.sqrt(q_base_sq_norm)) - q_rel = q_follower * wp.quat_inverse(q_rel_body) * wp.quat_inverse(q_base) - temp = X_T * R_base_T * quat_left_jacobian_inverse(q_rel) - if base_id_tot >= 0: - jac_rot_q_base = (-2.0 / q_base_sq_norm) * temp * G_of(q_base) - jac_rot_q_follower = (2.0 / q_follower_sq_norm) * temp * G_of(q_follower) + jac_rot_q_base, jac_rot_q_follower = _eval_rotation_jacobian_blocks( + X_T, q_base, q_follower, q_rel_body, base_id_tot >= 0 + ) # Note: we need X^T * R_base^T both for translation and rotation constraints, but to get the correct # derivatives for non-unit quaternions (which may be encountered before convergence) we end up needing # to use a separate formula to evaluate R_base in either case @@ -1100,13 +1271,9 @@ def _eval_joint_constraints_jacobian( return # Compute constraint Jacobian (cross product between x axis on base and y axis on follower) - a_x = X_T[0] - a_y = X_T[1] - if base_id_tot >= 0: - a_y_follower = unit_quat_apply(q_follower, a_y) - jac_q_base = -a_y_follower * unit_quat_apply_jacobian(q_base, a_x) - a_x_base = unit_quat_apply(q_base, a_x) - jac_q_follower = -a_x_base * unit_quat_apply_jacobian(q_follower, a_y) + jac_q_base, jac_q_follower = _eval_passive_universal_jacobian_blocks( + X_T, q_base, q_follower, base_id_tot >= 0 + ) # Write out Jacobian for i in range(4): @@ -1191,37 +1358,20 @@ def _eval_joint_constraints_sparse_jacobian( # Get base and follower transformations base_id = joints_bid_B[jt_id_tot] - if base_id < 0: - c_base = wp.vec3f(0.0, 0.0, 0.0) - q_base = wp.quatf(0.0, 0.0, 0.0, 1.0) - else: - c_base = wp.transform_get_translation(bodies_q[base_id]) - q_base = wp.transform_get_rotation(bodies_q[base_id]) follower_id = joints_bid_F[jt_id_tot] - c_follower = wp.transform_get_translation(bodies_q[follower_id]) - q_follower = wp.transform_get_rotation(bodies_q[follower_id]) + c_base, q_base, c_follower, q_follower = _load_joint_poses(base_id, follower_id, bodies_q) # Get target relative transformation (rotation part only, as translation part doesn't affect the Jacobian) q_rel_body = wp.transform_get_rotation(target_rel_transforms[jt_id_tot]) # Translation constraints - X_T_R_base_T = X_T * unit_quat_conj_to_rotation_matrix(q_base) - if base_id >= 0: - jac_trans_c_base = -X_T_R_base_T - delta_pos = unit_quat_apply(q_follower, x_follower) + c_follower - c_base - jac_trans_q_base = X_T * unit_quat_conj_apply_jacobian(q_base, delta_pos) - jac_trans_c_follower = X_T_R_base_T - jac_trans_q_follower = X_T_R_base_T * unit_quat_apply_jacobian(q_follower, x_follower) - + jac_trans_c_base, jac_trans_q_base, jac_trans_c_follower, jac_trans_q_follower = ( + _eval_translation_jacobian_blocks(X_T, q_base, q_follower, x_follower, c_base, c_follower, base_id >= 0) + ) # Rotation constraints - q_base_sq_norm = wp.dot(q_base, q_base) - q_follower_sq_norm = wp.dot(q_follower, q_follower) - R_base_T = unit_quat_conj_to_rotation_matrix(q_base / wp.sqrt(q_base_sq_norm)) - q_rel = q_follower * wp.quat_inverse(q_rel_body) * wp.quat_inverse(q_base) - temp = X_T * R_base_T * quat_left_jacobian_inverse(q_rel) - if base_id >= 0: - jac_rot_q_base = (-2.0 / q_base_sq_norm) * temp * G_of(q_base) - jac_rot_q_follower = (2.0 / q_follower_sq_norm) * temp * G_of(q_follower) + jac_rot_q_base, jac_rot_q_follower = _eval_rotation_jacobian_blocks( + X_T, q_base, q_follower, q_rel_body, base_id >= 0 + ) # Note: we need X^T * R_base^T both for translation and rotation constraints, but to get the correct # derivatives for non-unit quaternions (which may be encountered before convergence) we end up needing # to use a separate formula to evaluate R_base in either case @@ -1262,13 +1412,9 @@ def _eval_joint_constraints_sparse_jacobian( return # Compute constraint Jacobian (cross product between x axis on base and y axis on follower) - a_x = X_T[0] - a_y = X_T[1] - if base_id >= 0: - a_y_follower = unit_quat_apply(q_follower, a_y) - jac_q_base = -a_y_follower * unit_quat_apply_jacobian(q_base, a_x) - a_x_base = unit_quat_apply(q_base, a_x) - jac_q_follower = -a_x_base * unit_quat_apply_jacobian(q_follower, a_y) + jac_q_base, jac_q_follower = _eval_passive_universal_jacobian_blocks( + X_T, q_base, q_follower, base_id >= 0 + ) # Write out Jacobian if base_id >= 0: diff --git a/newton/_src/solvers/kamino/_src/solvers/fk/solver.py b/newton/_src/solvers/kamino/_src/solvers/fk/solver.py index 2119b70020..1643b76328 100644 --- a/newton/_src/solvers/kamino/_src/solvers/fk/solver.py +++ b/newton/_src/solvers/kamino/_src/solvers/fk/solver.py @@ -14,10 +14,12 @@ import numpy as np import warp as wp +from ......sim import ModelFlags from ....config import ForwardKinematicsSolverConfig from ...core.joints import JointActuationType, JointDoFType from ...core.model import ModelKamino from ...core.types import assign_to_warp_int32_array, to_warp_int32_array, vec7f +from ...kinematics.resets import get_base_q_from_joint_q_and_body_q from ...linalg.blas import ( block_sparse_ATA_blockwise_3_4_inv_diagonal_2d, block_sparse_ATA_inv_diagonal_2d, @@ -32,6 +34,8 @@ from .kernels import ( _add_regularizer_to_diagonal, _apply_line_search_step, + _compute_fk_axis_joint_frames, + _compute_fk_joint_frames, _correct_actuator_coords, _correct_universal_constraint_velocities, _eval_actuator_coords, @@ -52,6 +56,7 @@ _newton_check, _reset_state, _reset_state_base_q, + _resolve_fk_actuation_types, _update_cg_tolerance_kernel, create_1d_tile_based_kernels, create_2d_tile_based_kernels, @@ -59,6 +64,7 @@ create_eval_joint_constraints_kernel, create_eval_joint_constraints_sparse_jacobian_kernel, create_eval_min_num_iterations_kernel, + validate_fk_actuation_updates, ) from .types import FKJointDoFType, ForwardKinematicsPreconditionerType, ForwardKinematicsStatus @@ -183,13 +189,32 @@ def finalize(self, model: ModelKamino | None = None, config: ForwardKinematicsSo first_joint_id_prev = np.concatenate(([0], num_joints_prev.cumsum())) # Index of first joint per world # Resolve custom actuation types - joints_act_type_prev = self.model.joints.act_type.numpy().copy() if self.model.joints.fk_act_flag is not None: - joints_fk_act_flag = self.model.joints.fk_act_flag.numpy() - mapping = np.array([-1, JointActuationType.PASSIVE, JointActuationType.FORCE]) - joints_fk_act_type = mapping[joints_fk_act_flag + 1] # Map 0/1 flags to enum constants - overwrite_mask = joints_fk_act_flag != -1 - joints_act_type_prev[overwrite_mask] = joints_fk_act_type[overwrite_mask] + fk_act_flag = self.model.joints.fk_act_flag.numpy() + invalid = np.flatnonzero((fk_act_flag < -1) | (fk_act_flag > 1)) + if invalid.size > 0: + joint = int(invalid[0]) + raise ValueError(f"Invalid FK actuation flag for joint {joint}: expected -1, 0, or 1") + resolved_act_type = wp.empty( + shape=self.model.size.sum_of_num_joints, + dtype=wp.int32, + device=self.device, + ) + if self.model.size.sum_of_num_joints > 0: + wp.launch( + _resolve_fk_actuation_types, + dim=self.model.size.sum_of_num_joints, + inputs=[ + self.model.joints.act_type, + self.model.joints.fk_act_flag, + resolved_act_type, + ], + device=self.device, + ) + joints_act_type_prev = resolved_act_type.numpy() + # Indexed by model joint: 0 is passive, 1 is actuated, and -1 skips + # validation for an explicit base joint that FK replaces. + built_fk_actuated = (joints_act_type_prev != JointActuationType.PASSIVE).astype(np.int32) # Retrieve / compute dimensions - Actuated coordinates/dofs (main model) if self.model.joints.fk_act_flag is None: @@ -208,37 +233,38 @@ def finalize(self, model: ModelKamino | None = None, config: ForwardKinematicsSo classes = compute_fk_equivalence_classes(self.model) num_classes = len(classes) + # Resolve discrete joint data (e.g. types and indices) first, then + # copy or compute continuous joint data (e.g. frames). # Create a copy of the model's joints with added joints as needed: # - actuated free joints to reset the base position/orientation # - axis joints to factor out superfluous DoFs at tie rods joints_dof_type_prev = self.model.joints.dof_type.numpy().copy() joints_bid_B_prev = self.model.joints.bid_B.numpy().copy() joints_bid_F_prev = self.model.joints.bid_F.numpy().copy() - joints_B_r_Bj_prev = self.model.joints.B_r_Bj.numpy().copy() - joints_F_r_Fj_prev = self.model.joints.F_r_Fj.numpy().copy() - joints_X_Bj_prev = self.model.joints.X_Bj.numpy().copy() - joints_X_Fj_prev = self.model.joints.X_Fj.numpy().copy() joints_num_coords_prev = self.model.joints.num_coords.numpy().copy() joints_num_dofs_prev = self.model.joints.num_dofs.numpy().copy() joints_dof_type = [] joints_act_type = [] joints_bid_B = [] joints_bid_F = [] - joints_B_r_Bj = [] - joints_F_r_Fj = [] - joints_X_Bj = [] - joints_X_Fj = [] joints_num_actuated_coords = [] # Number of actuated coordinates per joint (0 for passive joints) joints_num_actuated_dofs = [] # Number of actuated dofs per joint (0 for passive joints) + joints_source_id = [] # Source joint in the main model, or -1 for synthetic joints + fk_axis_joint = [] # FK index of each synthetic axis joint + fk_axis_body = [] # Body defining each synthetic axis joint + fk_axis_source_joint_0 = [] # First source joint defining each synthetic axis joint + fk_axis_source_joint_1 = [] # Second source joint defining each synthetic axis joint num_joints = np.zeros(self.num_worlds, dtype=np.int32) # Number of joints per world self.num_joints_tot = 0 # Number of joints for all worlds actuated_coords_map = [] # Map of new actuated coordinates to these in the model or to the base coordinates actuated_dofs_map = [] # Map of new actuated dofs to these in the model or to the base dofs - base_q_default = np.zeros(7 * self.num_worlds, dtype=np.float32) # Default base pose - bodies_q_0 = self.model.bodies.q_i_0.numpy() base_joint_ids = self.num_worlds * [-1] # Base joint id per world base_joint_ids_input = self.model.info.base_joint_index.numpy().tolist() base_body_ids_input = self.model.info.base_body_index.numpy().tolist() + for base_joint_id in base_joint_ids_input: + if base_joint_id >= 0: + # FK always replaces an explicit base joint with an actuated free joint. + built_fk_actuated[base_joint_id] = -1 for wd_id in range(self.num_worlds): # Retrieve base joint id base_joint_id = base_joint_ids_input[wd_id] @@ -254,10 +280,7 @@ def finalize(self, model: ModelKamino | None = None, config: ForwardKinematicsSo joints_act_type.append(joints_act_type_prev[jt_id_prev]) joints_bid_B.append(joints_bid_B_prev[jt_id_prev]) joints_bid_F.append(joints_bid_F_prev[jt_id_prev]) - joints_B_r_Bj.append(joints_B_r_Bj_prev[jt_id_prev]) - joints_F_r_Fj.append(joints_F_r_Fj_prev[jt_id_prev]) - joints_X_Bj.append(joints_X_Bj_prev[jt_id_prev]) - joints_X_Fj.append(joints_X_Fj_prev[jt_id_prev]) + joints_source_id.append(jt_id_prev) if joints_act_type[-1] != JointActuationType.PASSIVE: num_coords_jt = joints_num_coords_prev[jt_id_prev] joints_num_actuated_coords.append(num_coords_jt) @@ -299,73 +322,24 @@ def finalize(self, model: ModelKamino | None = None, config: ForwardKinematicsSo joints_act_type.append(JointActuationType.PASSIVE) joints_bid_B.append(-1) joints_bid_F.append(rb_id_tot) - joints_B_r_Bj.append(np.zeros(dtype=np.float32, shape=3)) - joints_F_r_Fj.append(np.zeros(dtype=np.float32, shape=3)) + joints_source_id.append(-1) + fk_axis_joint.append(len(joints_dof_type) - 1) + fk_axis_body.append(rb_id_tot) + fk_axis_source_joint_0.append(spherical_joints_per_body[rb_id][0]) + fk_axis_source_joint_1.append(spherical_joints_per_body[rb_id][1]) joints_num_actuated_coords.append(0) joints_num_actuated_dofs.append(0) - # Compute position of both spherical joints on initial pose - def eval_joint_pos_init(jt_id_prev): - bid_B = joints_bid_B_prev[jt_id_prev] - bid_F = joints_bid_F_prev[jt_id_prev] - if bid_B == rb_id_tot: # Body is the joint's base # noqa: B023 - q_B = bodies_q_0[bid_B] - B_r_B = joints_B_r_Bj_prev[jt_id_prev] - return q_B[:3] + np.array(wp.quat_rotate(wp.quat(q_B[3:]), wp.vec3f(B_r_B))) - else: # Body is the joint's follower - assert bid_F == rb_id_tot # noqa: B023 - q_F = bodies_q_0[bid_F] - F_r_F = joints_F_r_Fj_prev[jt_id_prev] - return q_F[:3] + np.array(wp.quat_rotate(wp.quat(q_F[3:]), wp.vec3f(F_r_F))) - - pos_0 = eval_joint_pos_init(spherical_joints_per_body[rb_id][0]) - pos_1 = eval_joint_pos_init(spherical_joints_per_body[rb_id][1]) - - # Joint frame on base = in world coordinates - # Set X axis that connects both spherical joints (= tie rod axis) - a_x = pos_1 - pos_0 - a_x /= np.linalg.norm(a_x) - if np.abs(a_x[2]) < 0.99: - a_y = np.cross(np.array([0.0, 0.0, 1.0]), a_x) - else: - a_y = np.cross(np.array([0.0, 1.0, 0.0]), a_x) - a_y /= np.linalg.norm(a_y) - a_z = np.cross(a_x, a_y) - a_z /= np.linalg.norm(a_z) - axis_X_j = np.stack((a_x, a_y, a_z), axis=1) - joints_X_Bj.append(axis_X_j) - - # Joint frame on follower: set so that matches with frame on base on initial pose - q_F_0 = bodies_q_0[rb_id_tot][3:] - if np.max(np.abs(q_F_0 - np.array([0.0, 0.0, 0.0, 1.0]))) > 1e-8: - R_F_0_wp = wp.quat_to_matrix(wp.quatf(q_F_0)) - R_F_0 = np.reshape(np.array(R_F_0_wp), shape=(3, 3)) - joints_X_Fj.append(R_F_0 @ axis_X_j) - else: - joints_X_Fj.append(axis_X_j) - # Add joint for base joint / base body if base_joint_id >= 0: # Replace base joint with an actuated free joint joints_dof_type.append(JointDoFType.FREE) joints_act_type.append(JointActuationType.FORCE) joints_bid_B.append(-1) joints_bid_F.append(joints_bid_F_prev[base_joint_id]) - joints_B_r_Bj.append(joints_B_r_Bj_prev[base_joint_id]) - joints_F_r_Fj.append(joints_F_r_Fj_prev[base_joint_id]) - joints_X_Bj.append(joints_X_Bj_prev[base_joint_id]) - joints_X_Fj.append(joints_X_Fj_prev[base_joint_id]) + joints_source_id.append(base_joint_id) joints_num_actuated_coords.append(7) coord_offset = -7 * wd_id - 1 # We encode offsets in base_q negatively with i -> -i - 1 actuated_coords_map.extend(range(coord_offset, coord_offset - 7, -1)) - base_q_default[7 * wd_id : 7 * wd_id + 7] = [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - ] # Default to zero of free joint joints_num_actuated_dofs.append(6) dof_offset = -6 * wd_id - 1 # We encode offsets in base_u negatively with i -> -i - 1 actuated_dofs_map.extend(range(dof_offset, dof_offset - 6, -1)) @@ -376,17 +350,10 @@ def eval_joint_pos_init(jt_id_prev): joints_act_type.append(JointActuationType.FORCE) joints_bid_B.append(-1) joints_bid_F.append(base_body_id) - joints_B_r_Bj.append(np.zeros(3, dtype=np.float32)) - joints_F_r_Fj.append(np.zeros(3, dtype=np.float32)) - joints_X_Bj.append(np.eye(3, 3, dtype=np.float32)) - joints_X_Fj.append(np.eye(3, 3, dtype=np.float32)) + joints_source_id.append(-1) joints_num_actuated_coords.append(7) - # Note: we rely on the initial body orientations being identity - # Only then will the corresponding joint coordinates be interpretable as - # specifying the absolute base position and orientation coord_offset = -7 * wd_id - 1 # We encode offsets in base_q negatively with i -> -i - 1 actuated_coords_map.extend(range(coord_offset, coord_offset - 7, -1)) - base_q_default[7 * wd_id : 7 * wd_id + 7] = bodies_q_0[base_body_id] # Default to initial body pose joints_num_actuated_dofs.append(6) dof_offset = -6 * wd_id - 1 # We encode offsets in base_u negatively with i -> -i - 1 actuated_dofs_map.extend(range(dof_offset, dof_offset - 6, -1)) @@ -573,19 +540,29 @@ def eval_joint_pos_init(jt_id_prev): self.num_constraints = to_warp_int32_array(num_constraints) self.constraint_full_to_red_map = to_warp_int32_array(constraint_full_to_red_map) + # Helper data for model updates validation + self._built_fk_actuated = to_warp_int32_array(built_fk_actuated) + self._fk_actuation_violations = wp.empty(2, dtype=wp.int32) + # Modified joints self.joints_dof_type = to_warp_int32_array(joints_dof_type) self.joints_act_type = to_warp_int32_array(joints_act_type) self.joints_bid_B = to_warp_int32_array(joints_bid_B) self.joints_bid_F = to_warp_int32_array(joints_bid_F) - self.joints_B_r_Bj = wp.from_numpy(joints_B_r_Bj, dtype=wp.vec3f) - self.joints_F_r_Fj = wp.from_numpy(joints_F_r_Fj, dtype=wp.vec3f) - self.joints_X_Bj = wp.from_numpy(joints_X_Bj, dtype=wp.mat33f) - self.joints_X_Fj = wp.from_numpy(joints_X_Fj, dtype=wp.mat33f) + self.joints_B_r_Bj = wp.empty(self.num_joints_tot, dtype=wp.vec3f) + self.joints_F_r_Fj = wp.empty(self.num_joints_tot, dtype=wp.vec3f) + self.joints_X_Bj = wp.empty(self.num_joints_tot, dtype=wp.mat33f) + self.joints_X_Fj = wp.empty(self.num_joints_tot, dtype=wp.mat33f) + self.joints_source_id = to_warp_int32_array(joints_source_id) + self.fk_axis_joint = to_warp_int32_array(fk_axis_joint) + self.fk_axis_body = to_warp_int32_array(fk_axis_body) + self.fk_axis_source_joint_0 = to_warp_int32_array(fk_axis_source_joint_0) + self.fk_axis_source_joint_1 = to_warp_int32_array(fk_axis_source_joint_1) + self.num_axis_joints = len(fk_axis_joint) self.base_joint_id = to_warp_int32_array(base_joint_ids) # Default base state - self.base_q_default = wp.from_numpy(base_q_default, dtype=wp.transformf) + self.base_q_default = wp.zeros(shape=self.num_worlds, dtype=wp.transformf) self.base_u_default = wp.zeros(shape=(self.num_worlds,), dtype=wp.spatial_vectorf) # Line search @@ -931,6 +908,121 @@ def _cg_matvec_flat(x, y, world_active): maxiter=self.cg_max_iter, ) + # Initialize continuous joint data (e.g. joint frames) + self._update_joint_frames() + self._update_axis_joint_frames() + self._update_base_q_default() + + def validate_model_changed(self, flags: ModelFlags | int) -> None: + """Validate FK structural invariants before model values are updated. + + Args: + flags: Bitmask indicating which model properties changed. + + Raises: + RuntimeError: If the effective set of joints that are actuated for FK changed. + ValueError: If an FK actuation override is invalid. + """ + if not flags & (ModelFlags.JOINT_DOF_PROPERTIES | ModelFlags.ACTUATOR_PROPERTIES): + return + joint_count = self.model.size.sum_of_num_joints + if joint_count == 0: + return + + self._fk_actuation_violations.fill_(joint_count) + wp.launch( + validate_fk_actuation_updates, + dim=joint_count, + inputs=[ + self.model.joints.act_type, + self.model.joints.fk_act_flag, + self._built_fk_actuated, + self._fk_actuation_violations, + ], + device=self.device, + ) + changed_joint, invalid_joint = self._fk_actuation_violations.numpy() + if invalid_joint != joint_count: + raise ValueError(f"Invalid FK actuation flag for joint {int(invalid_joint)}: expected -1, 0, or 1") + if changed_joint != joint_count: + raise RuntimeError( + f"Changing the actuated vs passive status of joint {int(changed_joint)} for FK is not supported; " + "recreate SolverKamino to apply the change." + ) + + def notify_model_changed(self, flags: ModelFlags | int) -> None: + """Refresh FK-owned values after an in-place model update. + + Structural changes must be rejected by the owning solver before this + method is called. Updates here preserve allocations and pointers. + + Args: + flags: Bitmask indicating which model properties changed. + """ + if flags & (ModelFlags.JOINT_PROPERTIES | ModelFlags.BODY_INERTIAL_PROPERTIES): + self._update_joint_frames() + + if flags & (ModelFlags.JOINT_PROPERTIES | ModelFlags.BODY_PROPERTIES | ModelFlags.BODY_INERTIAL_PROPERTIES): + self._update_axis_joint_frames() + + if flags & (ModelFlags.JOINT_PROPERTIES | ModelFlags.BODY_PROPERTIES | ModelFlags.BODY_INERTIAL_PROPERTIES): + self._update_base_q_default() + + def _update_joint_frames(self) -> None: + """Compute FK joint frames from the current Kamino model.""" + if self.num_joints_tot == 0: + return + wp.launch( + _compute_fk_joint_frames, + dim=self.num_joints_tot, + inputs=[ + self.joints_source_id, + self.model.joints.B_r_Bj, + self.model.joints.F_r_Fj, + self.model.joints.X_Bj, + self.model.joints.X_Fj, + self.joints_B_r_Bj, + self.joints_F_r_Fj, + self.joints_X_Bj, + self.joints_X_Fj, + ], + device=self.device, + ) + + def _update_axis_joint_frames(self) -> None: + """Compute synthetic axis-joint frames from the current model.""" + if self.num_axis_joints == 0: + return + wp.launch( + _compute_fk_axis_joint_frames, + dim=self.num_axis_joints, + inputs=[ + self.fk_axis_joint, + self.fk_axis_body, + self.fk_axis_source_joint_0, + self.fk_axis_source_joint_1, + self.model.joints.bid_B, + self.model.joints.B_r_Bj, + self.model.joints.F_r_Fj, + self.model.bodies.q_i_0, + self.joints_X_Bj, + self.joints_X_Fj, + ], + device=self.device, + ) + + def _update_base_q_default(self) -> None: + """Compute default FK base poses from the current reference pose.""" + if self.num_worlds == 0: + return + get_base_q_from_joint_q_and_body_q( + model=self.model, + joint_q=self.model.joints.q_j_0, + body_q=self.model.bodies.q_i_0, + base_q=self.base_q_default, + world_mask=self.all_worlds_mask, + ) + ### # Internal evaluators (graph-capturable functions working on pre-allocated data) ### diff --git a/newton/_src/solvers/kamino/_src/solvers/metrics.py b/newton/_src/solvers/kamino/_src/solvers/metrics.py index c1457ee95d..edfd29f72e 100644 --- a/newton/_src/solvers/kamino/_src/solvers/metrics.py +++ b/newton/_src/solvers/kamino/_src/solvers/metrics.py @@ -79,7 +79,6 @@ import warp as wp from ..core.data import DataKamino -from ..core.math import screw, screw_angular, screw_linear from ..core.model import ModelKamino from ..core.state import StateKamino from ..core.types import vec6f @@ -241,10 +240,9 @@ class SolutionMetricsData: """ The largest constraint violation residual across all contact constraints. - Computed as the maximum absolute value (i.e. infinity-norm) over contact constraint residuals. - - Equivalent to `r_cts_contacts := || d_k ||_inf`, where `d_k` would be an array of - contact penetrations extracted from the `gapfunc` elements of :class:`ContactsKaminoData`. + Equivalent to `r_cts_contacts := max_k max(0, -d_k)`, where `d_k` is the + margin-shifted signed distance stored in the ``w`` component of the contact + `gapfunc`. Negative `d_k` denotes penetration. Shape of ``(num_worlds,)``. """ @@ -592,7 +590,7 @@ def compute_vector_difference_infnorm( def _compute_eom_residual( # Inputs model_time_dt: wp.array[wp.float32], - model_gravity: wp.array[wp.vec4f], + model_gravity: wp.array[wp.vec3f], model_bodies_wid: wp.array[wp.int32], model_bodies_m_i: wp.array[wp.float32], state_bodies_I_i: wp.array[wp.mat33f], @@ -616,22 +614,21 @@ def _compute_eom_residual( # Retrieve the time step dt = model_time_dt[wid] - gravity = model_gravity[wid] - g = gravity.w * wp.vec3f(gravity.x, gravity.y, gravity.z) + g = model_gravity[wid] # Decompose into linear and angular parts - f_i = screw_linear(w_i) - v_i = screw_linear(u_i) - v_i_p = screw_linear(u_i_p) - tau_i = screw_angular(w_i) - omega_i = screw_angular(u_i) - omega_i_p = screw_angular(u_i_p) + f_i = wp.spatial_top(w_i) + v_i = wp.spatial_top(u_i) + v_i_p = wp.spatial_top(u_i_p) + tau_i = wp.spatial_bottom(w_i) + omega_i = wp.spatial_bottom(u_i) + omega_i_p = wp.spatial_bottom(u_i_p) S_i = wp.skew(omega_i_p) # Compute the per-body EoM residual over linear and angular parts r_linear_i = wp.abs(m_i * (v_i - v_i_p) - dt * (m_i * g + f_i)) r_angular_i = wp.abs(I_i @ (omega_i - omega_i_p) - dt * (tau_i - S_i @ (I_i @ omega_i_p))) - r_i = screw(r_linear_i, r_angular_i) + r_i = wp.spatial_vectorf(*r_linear_i, *r_angular_i) # Compute the per-body maximum residual and argmax index r_eom_i = wp.max(r_i) @@ -873,12 +870,12 @@ def _compute_cts_contacts_residual( wcid = contact_cid[cid] gapfunc = contact_gapfunc[cid] - # Compute the per-contact constraint residual (infinity-norm) - r_cts_contacts_k = wp.abs(gapfunc[3]) + # Compute unilateral penetration depth from the margin-shifted signed distance. + r_cts_contacts_k = wp.max(0.0, -gapfunc[3]) # Update the per-world maximum residual and argmax index previous_max = wp.atomic_max(metric_r_cts_contacts, wid, r_cts_contacts_k) - if r_cts_contacts_k >= previous_max: + if r_cts_contacts_k > 0.0 and r_cts_contacts_k >= previous_max: wp.atomic_exch(metric_r_cts_contacts_argmax, wid, wcid) diff --git a/newton/_src/solvers/kamino/_src/solvers/padmm/kernels.py b/newton/_src/solvers/kamino/_src/solvers/padmm/kernels.py index f154140ec8..69b1db7481 100644 --- a/newton/_src/solvers/kamino/_src/solvers/padmm/kernels.py +++ b/newton/_src/solvers/kamino/_src/solvers/padmm/kernels.py @@ -11,6 +11,18 @@ import warp as wp from ...core.math import FLOAT32_EPS, FLOAT32_MAX +from ..common import ( + apply_dual_preconditioner_to_solution as _apply_dual_preconditioner_to_solution, +) +from ..common import ( + warmstart_contact_constraints as _warmstart_contact_constraints, +) +from ..common import ( + warmstart_joint_constraints as _warmstart_joint_constraints, +) +from ..common import ( + warmstart_limit_constraints as _warmstart_limit_constraints, +) from .math import ( compute_cwise_vec_div, compute_cwise_vec_mul, @@ -151,187 +163,6 @@ def _warmstart_desaxce_correction( solver_z[ccio_k + 2] = vn + mu * vt_norm -@wp.kernel -def _warmstart_joint_constraints( - # Inputs: - model_time_dt: wp.array[wp.float32], - joint_wid: wp.array[wp.int32], - joint_num_dynamic_cts: wp.array[wp.int32], - joint_num_kinematic_cts: wp.array[wp.int32], - joint_dynamic_cts_offset_joint_cts: wp.array[wp.int32], - joint_kinematic_cts_offset_joint_cts: wp.array[wp.int32], - joint_dynamic_cts_offset_total_cts: wp.array[wp.int32], - joint_kinematic_cts_offset_total_cts: wp.array[wp.int32], - joint_lambda_j: wp.array[wp.float32], - problem_P: wp.array[wp.float32], - # Outputs: - x_0: wp.array[wp.float32], - y_0: wp.array[wp.float32], - z_0: wp.array[wp.float32], -): - # Retrieve the thread index as the joint index - jid = wp.tid() - - # Retrieve the joint-specific model info - wid_j = joint_wid[jid] - num_dynamic_cts_j = joint_num_dynamic_cts[jid] - num_kinematic_cts_j = joint_num_kinematic_cts[jid] - - # Retrieve the world-specific info - dt = model_time_dt[wid_j] - - # Retrieve offsets in the joint-only and total constraints vector - joint_dyn_cts_start = joint_dynamic_cts_offset_joint_cts[jid] - joint_kin_cts_start = joint_kinematic_cts_offset_joint_cts[jid] - dyn_cts_row_start_j = joint_dynamic_cts_offset_total_cts[jid] - kin_cts_row_start_j = joint_kinematic_cts_offset_total_cts[jid] - - # For each joint constraint, scale the constraint force by the time-step and - # the preconditioner and initialize the solver state variables accordingly - for j in range(num_dynamic_cts_j): - P_j = problem_P[dyn_cts_row_start_j + j] - lambda_j = (dt / P_j) * joint_lambda_j[joint_dyn_cts_start + j] - x_0[dyn_cts_row_start_j + j] = lambda_j - y_0[dyn_cts_row_start_j + j] = lambda_j - z_0[dyn_cts_row_start_j + j] = 0.0 - for j in range(num_kinematic_cts_j): - P_j = problem_P[kin_cts_row_start_j + j] - lambda_j = (dt / P_j) * joint_lambda_j[joint_kin_cts_start + j] - x_0[kin_cts_row_start_j + j] = lambda_j - y_0[kin_cts_row_start_j + j] = lambda_j - z_0[kin_cts_row_start_j + j] = 0.0 - - -@wp.kernel -def _warmstart_limit_constraints( - # Inputs: - model_time_dt: wp.array[wp.float32], - model_info_total_cts_offset: wp.array[wp.int32], - data_info_limit_cts_group_offset: wp.array[wp.int32], - limit_model_num_active: wp.array[wp.int32], - limit_wid: wp.array[wp.int32], - limit_lid: wp.array[wp.int32], - limit_reaction: wp.array[wp.float32], - limit_velocity: wp.array[wp.float32], - problem_P: wp.array[wp.float32], - # Outputs: - x_0: wp.array[wp.float32], - y_0: wp.array[wp.float32], - z_0: wp.array[wp.float32], -): - # Retrieve the thread index as the limit index - lid = wp.tid() - - # Retrieve the number of limits active in the model - model_nl = limit_model_num_active[0] - - # Skip if lid is greater than the number of limits active in the model - if lid >= model_nl: - return - - # Retrieve the limit-specific data - wid = limit_wid[lid] - lid_l = limit_lid[lid] - lambda_l = limit_reaction[lid] - v_plus_l = limit_velocity[lid] - - # Retrieve the world-specific info - dt = model_time_dt[wid] - total_cts_offset = model_info_total_cts_offset[wid] - limit_cts_offset = data_info_limit_cts_group_offset[wid] - - # Compute block offsets of the limit constraints within - # the limit-only constraints and total constraints arrays - vio_l = total_cts_offset + limit_cts_offset + lid_l - - # Load the diagonal preconditioner for the limit constraints - # NOTE: We only need to load the first element since by necessity - # the preconditioner is constant across the 3 constraint dimensions - P_l = problem_P[vio_l] - - # Scale the limit force by the time-step to - # render an impulse and by the preconditioner - lambda_l *= dt / P_l - - # Scale the limit velocity by the preconditioner - v_plus_l *= P_l - - # Compute and store the limit-constraint reaction forces - x_0[vio_l] = lambda_l - y_0[vio_l] = lambda_l - z_0[vio_l] = v_plus_l - - -@wp.kernel -def _warmstart_contact_constraints( - # Inputs: - model_time_dt: wp.array[wp.float32], - model_info_total_cts_offset: wp.array[wp.int32], - data_info_contact_cts_group_offset: wp.array[wp.int32], - contact_model_num_contacts: wp.array[wp.int32], - contact_wid: wp.array[wp.int32], - contact_cid: wp.array[wp.int32], - contact_material: wp.array[wp.vec2f], - contact_reaction: wp.array[wp.vec3f], - contact_velocity: wp.array[wp.vec3f], - problem_P: wp.array[wp.float32], - # Outputs: - x_0: wp.array[wp.float32], - y_0: wp.array[wp.float32], - z_0: wp.array[wp.float32], -): - # Retrieve the thread index as the contact index - cid = wp.tid() - - # Retrieve the number of contacts active in the model - model_nc = contact_model_num_contacts[0] - - # Skip if cid is greater than the number of contacts active in the model - if cid >= model_nc: - return - - # Retrieve the contact-specific data - wid = contact_wid[cid] - cid_k = contact_cid[cid] - material_k = contact_material[cid] - lambda_k = contact_reaction[cid] - v_plus_k = contact_velocity[cid] - - # Retrieve the world-specific info - dt = model_time_dt[wid] - total_cts_offset = model_info_total_cts_offset[wid] - contact_cts_offset = data_info_contact_cts_group_offset[wid] - - # Compute block offsets of the contact constraints within - # the contact-only constraints and total constraints arrays - vio_k = total_cts_offset + contact_cts_offset + 3 * cid_k - - # Load the diagonal preconditioner for the contact constraints - # NOTE: We only need to load the first element since by necessity - # the preconditioner is constant across the 3 constraint dimensions - P_k = problem_P[vio_k] - - # Scale the contact force by the time-step to - # render an impulse and by the preconditioner - lambda_k *= dt / P_k - - # Scale the contact velocity by the preconditioner - # and apply the De Saxce correction to the post-event - # contact velocity to render solver dual variables - v_plus_k *= P_k - mu_k = material_k[0] - vt_norm = wp.sqrt(v_plus_k.x * v_plus_k.x + v_plus_k.y * v_plus_k.y) - v_plus_k.z += mu_k * vt_norm - - # Compute and store the contact-constraint reaction forces - for k in range(3): - x_0[vio_k + k] = lambda_k[k] - for k in range(3): - y_0[vio_k + k] = lambda_k[k] - for k in range(3): - z_0[vio_k + k] = v_plus_k[k] - - def make_initialize_solver_kernel(use_acceleration: bool = False): """ Creates a kernel to initialize the PADMM solver state, status, and penalty parameters. @@ -1823,44 +1654,6 @@ def _apply_dual_preconditioner_to_state( solver_z[v_i] = (1.0 / P_i) * z_i -@wp.kernel -def _apply_dual_preconditioner_to_solution( - # Inputs: - problem_dim: wp.array[wp.int32], - problem_vio: wp.array[wp.int32], - problem_P: wp.array[wp.float32], - # Outputs: - solution_lambdas: wp.array[wp.float32], - solution_v_plus: wp.array[wp.float32], -): - # Retrieve the thread index - wid, tid = wp.tid() - - # Retrieve the number of active constraints in the world - ncts = problem_dim[wid] - - # Skip if row index exceed the problem size - if tid >= ncts: - return - - # Retrieve the vector index offset of the world - vio = problem_vio[wid] - - # Compute the global index of the vector entry - v_i = vio + tid - - # Retrieve the i-th entries of the target vectors - lambdas_i = solution_lambdas[v_i] - v_plus_i = solution_v_plus[v_i] - - # Retrieve the i-th entry of the diagonal preconditioner - P_i = problem_P[v_i] - - # Store the preconditioned i-th entry of the vector - solution_lambdas[v_i] = (1.0 / P_i) * lambdas_i - solution_v_plus[v_i] = P_i * v_plus_i - - @wp.kernel def _compute_final_desaxce_correction( problem_nc: wp.array[wp.int32], diff --git a/newton/_src/solvers/kamino/_src/solvers/padmm/types.py b/newton/_src/solvers/kamino/_src/solvers/padmm/types.py index 223abbf864..31e7cbb165 100644 --- a/newton/_src/solvers/kamino/_src/solvers/padmm/types.py +++ b/newton/_src/solvers/kamino/_src/solvers/padmm/types.py @@ -34,7 +34,6 @@ from __future__ import annotations from enum import IntEnum -from typing import Any import numpy as np import warp as wp @@ -43,6 +42,7 @@ from ....config import PADMMSolverConfig from ...core.size import SizeKamino from ...core.types import to_warp_int32_array +from ..common import DualSolution, WarmStartMode ### # Module interface @@ -126,60 +126,8 @@ def __repr__(self): return self.__str__() -class PADMMWarmStartMode(IntEnum): - """ - An enumeration of the warmstart modes used in PADMM. - """ - - NONE = -1 - """ - No warmstart: - The solver does not use any warmstart information and starts from - scratch, i.e. performs a cold-start regardless of any cached state. - """ - - INTERNAL = 0 - """ - From internally cached solution: - The solver uses its values currently in the solution - container as warmstart information for the current solve. - """ - - CONTAINERS = 1 - """ - From externally cached solution containers: - The solver uses values from externally provided solution - containers as warmstart information for the current solve. - """ - - @classmethod - def from_string(cls, s: str) -> PADMMWarmStartMode: - """Converts a string to a PADMMWarmStartMode enum value.""" - try: - return cls[s.upper()] - except KeyError as e: - raise ValueError(f"Invalid PADMMWarmStartMode: {s}. Valid options are: {[e.name for e in cls]}") from e - - @override - def __str__(self): - """Returns a string representation of the PADMMWarmStartMode.""" - return f"PADMMWarmStartMode.{self.name} ({self.value})" - - @override - def __repr__(self): - """Returns a string representation of the PADMMWarmStartMode.""" - return self.__str__() - - @staticmethod - def parse_usd_attribute(value: str, context: dict[str, Any] | None = None) -> str: - """Parse warmstart option imported from USD, following the KaminoSceneAPI schema.""" - if not isinstance(value, str): - raise TypeError("Parser expects input of type 'str'.") - mapping = {"none": "none", "internal": "internal", "containers": "containers"} - lower_value = value.lower().strip() - if lower_value not in mapping: - raise ValueError(f"Warmstart parameter '{value}' is not a valid option.") - return mapping[lower_value] +PADMMWarmStartMode = WarmStartMode +"""Backward-compatible alias for the shared Kamino warm-start mode.""" @wp.struct @@ -815,59 +763,8 @@ def zero(self, use_acceleration: bool = False): self.r_dz.zero_() -class PADMMSolution: - """ - An interface container to the PADMM solver solution arrays. - - Attributes: - lambdas: The constraint reactions (i.e. impulses) solution array. - Shape of ``(sum_of_max_total_cts,)``. - v_plus: The post-event constraint-space velocities solution array. - Shape of ``(sum_of_max_total_cts,)``. - """ - - def __init__(self, size: SizeKamino | None = None): - """ - Initializes the PADMM solution container. - - If a model size is provided, allocates the solution arrays accordingly. - - Args: - size: The model-size utility container holding the dimensionality of the model. - """ - - self.lambdas: wp.array[wp.float32] | None = None - """ - The constraint reactions (i.e. impulses) solution array. - Shape of ``(sum_of_max_total_cts,)``. - """ - - self.v_plus: wp.array[wp.float32] | None = None - """ - The post-event constraint-space velocities solution array. - Shape of ``(sum_of_max_total_cts,)``. - """ - - # Perform memory allocations if model size is specified - if size is not None: - self.finalize(size) - - def finalize(self, size: SizeKamino): - """ - Allocates the PADMM solution arrays based on the model size. - - Args: - size: The model-size utility container holding the dimensionality of the model. - """ - self.lambdas = wp.zeros(size.sum_of_max_total_cts, dtype=wp.float32) - self.v_plus = wp.zeros(size.sum_of_max_total_cts, dtype=wp.float32) - - def zero(self): - """ - Resets all PADMM solution arrays to zeros. - """ - self.lambdas.zero_() - self.v_plus.zero_() +PADMMSolution = DualSolution +"""Backward-compatible alias for the shared dual-solver solution container.""" class PADMMInfo: diff --git a/newton/_src/solvers/kamino/_src/solvers/warmstart.py b/newton/_src/solvers/kamino/_src/solvers/warmstart.py index fc15984c89..b8fee1cb7c 100644 --- a/newton/_src/solvers/kamino/_src/solvers/warmstart.py +++ b/newton/_src/solvers/kamino/_src/solvers/warmstart.py @@ -225,6 +225,8 @@ def _warmstart_contacts_by_matched_geom_pair_key_and_position( # Update the current old-key to check in the next iteration k += 1 + if start + k >= num_active_old: + break old_key = sorted_contact_keys_old[start + k] # Store the new contact reaction and velocity @@ -340,6 +342,8 @@ def _warmstart_contacts_from_geom_pair_net_force( # Update the current old-key to check in the next iteration k += 1 + if start + k >= num_active_old: + break old_key = sorted_contact_keys_old[start + k] # TODO: We need to cache this value per geom-pair @@ -470,6 +474,8 @@ def _warmstart_contacts_by_matched_geom_pair_key_and_position_with_net_force_bac # Update the current old-key to check in the next iteration k += 1 + if start + k >= num_active_old: + break old_key = sorted_contact_keys_old[start + k] # If no matching contact found by position, fallback to net wrench approach @@ -519,6 +525,8 @@ def _warmstart_contacts_by_matched_geom_pair_key_and_position_with_net_force_bac # Update the current old-key to check in the next iteration k += 1 + if start + k >= num_active_old: + break old_key = sorted_contact_keys_old[start + k] # TODO: We need to cache this value per geom-pair diff --git a/newton/_src/solvers/kamino/_src/utils/benchmark/__main__.py b/newton/_src/solvers/kamino/_src/utils/benchmark/__main__.py index b7dc3ebc58..9dbe1655c6 100644 --- a/newton/_src/solvers/kamino/_src/utils/benchmark/__main__.py +++ b/newton/_src/solvers/kamino/_src/utils/benchmark/__main__.py @@ -226,7 +226,7 @@ def benchmark_run(args: argparse.Namespace): msg.notif("[Device]: %s", spec_info) # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"using_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/_src/utils/benchmark/problems.py b/newton/_src/solvers/kamino/_src/utils/benchmark/problems.py index 126c271134..68de1ce99a 100644 --- a/newton/_src/solvers/kamino/_src/utils/benchmark/problems.py +++ b/newton/_src/solvers/kamino/_src/utils/benchmark/problems.py @@ -99,8 +99,9 @@ def builder_fn(): build_fn=basics.build_boxes_fourbar, ground=ground, ) - for w in range(num_worlds): - builder.gravity[w].enabled = gravity + if not gravity: + for w in range(num_worlds): + builder.set_gravity(wp.vec3f(0.0), w) return builder control = ControlConfig(decimation=20, scale=10.0) @@ -142,8 +143,9 @@ def builder_fn(): for w in range(num_worlds): add_ground_box(builder, world_index=w) # Set gravity - for w in range(builder.num_worlds): - builder.gravity[w].enabled = gravity + if not gravity: + for w in range(builder.num_worlds): + builder.set_gravity(wp.vec3f(0.0), w) return builder # Set control configurations diff --git a/newton/_src/solvers/kamino/_src/utils/io/usd.py b/newton/_src/solvers/kamino/_src/utils/io/usd.py index 68fc9f9c1d..ebf10ce75a 100644 --- a/newton/_src/solvers/kamino/_src/utils/io/usd.py +++ b/newton/_src/solvers/kamino/_src/utils/io/usd.py @@ -3,6 +3,7 @@ """Provides mechanisms to import OpenUSD Physics models.""" +import math import uuid from collections.abc import Iterable from pathlib import Path @@ -28,13 +29,12 @@ JointDoFType, ) from ...core.materials import ( - DEFAULT_DENSITY, DEFAULT_FRICTION, DEFAULT_RESTITUTION, MaterialDescriptor, MaterialPairProperties, ) -from ...core.math import I_3, axis_to_mat33, screw +from ...core.math import I_3, axis_to_mat33 from ...core.shapes import ( BoxShape, CapsuleShape, @@ -501,12 +501,9 @@ def _parse_material( ### # Retrieve the USD material properties - density_scale = mass_unit / distance_unit**3 - density = (density_scale) * self._parse_float(material_prim, "physics:density", default=DEFAULT_DENSITY) restitution = self._parse_float(material_prim, "physics:restitution", default=DEFAULT_RESTITUTION) static_friction = self._parse_float(material_prim, "physics:staticFriction", default=DEFAULT_FRICTION) dynamic_friction = self._parse_float(material_prim, "physics:dynamicFriction", default=DEFAULT_FRICTION) - msg.debug(f"density: {density}") msg.debug(f"restitution: {restitution}") msg.debug(f"static_friction: {static_friction}") msg.debug(f"dynamic_friction: {dynamic_friction}") @@ -518,7 +515,6 @@ def _parse_material( return MaterialDescriptor( name=name, uid=uid, - density=density, restitution=restitution, static_friction=static_friction, dynamic_friction=dynamic_friction, @@ -660,7 +656,7 @@ def _parse_rigid_body( # Construct the initial pose and twist of the body in world coordinates q_i_0 = wp.transformf(r_com_i, body_xform.q) - u_i_0 = screw(v_i, omega_i) + u_i_0 = wp.spatial_vectorf(*v_i, *omega_i) msg.debug(f"q_i_0: {q_i_0}") msg.debug(f"u_i_0: {u_i_0}") @@ -708,7 +704,13 @@ def _make_joint_default_limits(self, dof_type: JointDoFType) -> tuple[list[float num_dofs = int(dof_type.num_dofs) q_j_min = [JOINT_QMIN] * num_dofs q_j_max = [JOINT_QMAX] * num_dofs - tau_j_max = [JOINT_TAUMAX] * num_dofs + # Mirroring Newton USD importer behavior: Default effort limit on + # revolute/prismatic/spherical joints is set to JOINT_TAUMAX. On all + # D6-derived joints, the effort limit is set to `math.inf`. + if dof_type in (JointDoFType.REVOLUTE, JointDoFType.PRISMATIC, JointDoFType.SPHERICAL): + tau_j_max = [JOINT_TAUMAX] * num_dofs + else: + tau_j_max = [math.inf] * num_dofs return q_j_min, q_j_max, tau_j_max def _make_joint_default_dynamics( @@ -748,7 +750,8 @@ def _parse_joint_revolute( q_j_max[0] = min(rotation_unit * joint_spec.limit.upper, JOINT_QMAX) if joint_spec.drive.enabled: if not joint_spec.drive.acceleration: - tau_j_max[0] = min(joint_spec.drive.forceLimit, JOINT_TAUMAX) + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. + tau_j_max[0] = joint_spec.drive.forceLimit has_pd_gains = joint_spec.drive.stiffness > 0.0 or joint_spec.drive.damping > 0.0 if load_drive_dynamics and has_pd_gains: a_j = [0.0] * dof_type.num_dofs @@ -778,7 +781,8 @@ def _parse_joint_prismatic(self, joint_spec, distance_unit: float = 1.0, load_dr q_j_max[0] = min(distance_unit * joint_spec.limit.upper, JOINT_QMAX) if joint_spec.drive.enabled: if not joint_spec.drive.acceleration: - tau_j_max[0] = min(joint_spec.drive.forceLimit, JOINT_TAUMAX) + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. + tau_j_max[0] = joint_spec.drive.forceLimit has_pd_gains = joint_spec.drive.stiffness > 0.0 or joint_spec.drive.damping > 0.0 if load_drive_dynamics and has_pd_gains: a_j = [0.0] * dof_type.num_dofs @@ -815,7 +819,8 @@ def _parse_joint_revolute_from_d6(self, name, joint_prim, joint_spec, joint_dof, act_type = JointActuationType.FORCE for drive in joint_spec.jointDrives: if drive.first == joint_dof: - tau_j_max[0] = min(drive.second.forceLimit, JOINT_TAUMAX) + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. + tau_j_max[0] = drive.second.forceLimit else: act_type = JointActuationType.PASSIVE return dof_type, act_type, X_j, q_j_min, q_j_max, tau_j_max @@ -839,7 +844,8 @@ def _parse_joint_prismatic_from_d6(self, name, joint_prim, joint_spec, joint_dof act_type = JointActuationType.FORCE for drive in joint_spec.jointDrives: if drive.first == joint_dof: - tau_j_max[0] = min(drive.second.forceLimit, JOINT_TAUMAX) + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. + tau_j_max[0] = drive.second.forceLimit else: act_type = JointActuationType.PASSIVE return dof_type, act_type, X_j, q_j_min, q_j_max, tau_j_max @@ -867,10 +873,11 @@ def _parse_joint_cylindrical_from_d6( act_type = JointActuationType.FORCE for drive in joint_spec.jointDrives: dof = drive.first + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. if dof == self.UsdPhysics.JointDOF.TransX: - tau_j_max[0] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[0] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.RotX: - tau_j_max[1] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[1] = drive.second.forceLimit else: act_type = JointActuationType.PASSIVE return dof_type, act_type, q_j_min, q_j_max, tau_j_max @@ -896,10 +903,11 @@ def _parse_joint_universal_from_d6(self, name, joint_prim, joint_spec, rotation_ act_type = JointActuationType.FORCE for drive in joint_spec.jointDrives: dof = drive.first + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. if dof == self.UsdPhysics.JointDOF.RotX: - tau_j_max[0] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[0] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.RotY: - tau_j_max[1] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[1] = drive.second.forceLimit else: act_type = JointActuationType.PASSIVE return dof_type, act_type, q_j_min, q_j_max, tau_j_max @@ -934,12 +942,13 @@ def _parse_joint_cartesian_from_d6( act_type = JointActuationType.FORCE for drive in joint_spec.jointDrives: dof = drive.first + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. if dof == self.UsdPhysics.JointDOF.TransX: - tau_j_max[0] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[0] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.TransY: - tau_j_max[1] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[1] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.TransZ: - tau_j_max[2] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[2] = drive.second.forceLimit else: act_type = JointActuationType.PASSIVE return dof_type, act_type, q_j_min, q_j_max, tau_j_max @@ -968,12 +977,13 @@ def _parse_joint_spherical_from_d6(self, name, joint_prim, joint_spec, rotation_ act_type = JointActuationType.FORCE for drive in joint_spec.jointDrives: dof = drive.first + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. if dof == self.UsdPhysics.JointDOF.RotX: - tau_j_max[0] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[0] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.RotY: - tau_j_max[1] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[1] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.RotZ: - tau_j_max[2] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[2] = drive.second.forceLimit else: act_type = JointActuationType.PASSIVE return dof_type, act_type, q_j_min, q_j_max, tau_j_max @@ -1000,14 +1010,14 @@ def _parse_joint_free_from_d6( q_j_min[2] = max(distance_unit * limit.second.lower, JOINT_QMIN) q_j_max[2] = min(distance_unit * limit.second.upper, JOINT_QMAX) elif dof == self.UsdPhysics.JointDOF.RotX: - q_j_min[0] = max(rotation_unit * limit.second.lower, JOINT_QMIN) - q_j_max[0] = min(rotation_unit * limit.second.upper, JOINT_QMAX) + q_j_min[3] = max(rotation_unit * limit.second.lower, JOINT_QMIN) + q_j_max[3] = min(rotation_unit * limit.second.upper, JOINT_QMAX) elif dof == self.UsdPhysics.JointDOF.RotY: - q_j_min[1] = max(rotation_unit * limit.second.lower, JOINT_QMIN) - q_j_max[1] = min(rotation_unit * limit.second.upper, JOINT_QMAX) + q_j_min[4] = max(rotation_unit * limit.second.lower, JOINT_QMIN) + q_j_max[4] = min(rotation_unit * limit.second.upper, JOINT_QMAX) elif dof == self.UsdPhysics.JointDOF.RotZ: - q_j_min[2] = max(rotation_unit * limit.second.lower, JOINT_QMIN) - q_j_max[2] = min(rotation_unit * limit.second.upper, JOINT_QMAX) + q_j_min[5] = max(rotation_unit * limit.second.lower, JOINT_QMIN) + q_j_max[5] = min(rotation_unit * limit.second.upper, JOINT_QMAX) num_drives = len(joint_spec.jointDrives) if num_drives > 0: @@ -1019,18 +1029,19 @@ def _parse_joint_free_from_d6( act_type = JointActuationType.FORCE for drive in joint_spec.jointDrives: dof = drive.first + # To align with the Newton USD importer, the effort limit is not clamped to JOINT_TAUMAX. if dof == self.UsdPhysics.JointDOF.TransX: - tau_j_max[0] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[0] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.TransY: - tau_j_max[1] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[1] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.TransZ: - tau_j_max[2] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[2] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.RotX: - tau_j_max[0] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[3] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.RotY: - tau_j_max[1] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[4] = drive.second.forceLimit elif dof == self.UsdPhysics.JointDOF.RotZ: - tau_j_max[2] = min(drive.second.forceLimit, JOINT_TAUMAX) + tau_j_max[5] = drive.second.forceLimit else: act_type = JointActuationType.PASSIVE return dof_type, act_type, q_j_min, q_j_max, tau_j_max @@ -1800,8 +1811,7 @@ def import_from( # World ### - # Initialize the world properties - gravity = GravityDescriptor() + stage_up_axis = Axis.from_string(str(self.UsdGeom.GetStageUpAxis(stage))) # Parse for PhysicsScene prims if self.UsdPhysics.ObjectType.Scene in ret_dict: @@ -1813,16 +1823,25 @@ def import_from( msg.error("Multiple PhysicsScene prims found in the USD file. Only the first prim will be considered.") # Extract the world gravity from the physics scene - gravity.acceleration = distance_unit * scene_desc.gravityMagnitude - gravity.direction = wp.vec3f(scene_desc.gravityDirection) + gravity = GravityDescriptor.from_usd( + scene_desc.gravityDirection, + scene_desc.gravityMagnitude, + stage_up_axis, + distance_unit, + ) builder.set_gravity(gravity) - msg.debug(f"World gravity: {gravity}") - - # Set the world up-axis based on the gravity direction - up_axis = Axis.from_any(int(np.argmax(np.abs(scene_desc.gravityDirection)))) + msg.debug(f"World gravity: {gravity.vector}") + + # Set the world up-axis based on the resolved gravity vector. + gravity_vector = np.asarray(gravity.vector, dtype=np.float32) + up_axis = ( + stage_up_axis + if np.linalg.norm(gravity_vector) == 0.0 + else Axis.from_any(int(np.argmax(np.abs(gravity_vector)))) + ) else: # NOTE: Gravity is left with default values - up_axis = Axis.from_string(str(self.UsdGeom.GetStageUpAxis(stage))) + up_axis = stage_up_axis # Determine the up-axis transformation if apply_up_axis_from_stage: diff --git a/newton/_src/solvers/kamino/config.py b/newton/_src/solvers/kamino/config.py index de3dce4b3f..b09fffd4a7 100644 --- a/newton/_src/solvers/kamino/config.py +++ b/newton/_src/solvers/kamino/config.py @@ -25,6 +25,7 @@ "ConfigBase", "ConstrainedDynamicsConfig", "ConstraintStabilizationConfig", + "DVISolverConfig", "ForwardKinematicsSolverConfig", "PADMMSolverConfig", ] @@ -99,9 +100,9 @@ class CollisionDetectorConfig(ConfigBase): max_contacts: int | None = None """ - The maximum number of contacts to generate over the entire model.\n - Used to compute the total maximum contacts allocated for the model, - in conjunction with the total number of candidate geom-pairs.\n + The maximum number of contacts to allocate over the entire model.\n + Caps the geometry-based contact capacity estimate during collision-detector + initialization.\n Defaults to `DEFAULT_MODEL_MAX_CONTACTS` (`1000`) if unspecified. """ @@ -774,6 +775,185 @@ def __post_init__(self): self.validate() +@dataclass +class DVISolverConfig: + """ + A container to hold configurations for the DVI forward dynamics solver. + """ + + max_iterations: int = 100 + """ + Maximum projected Gauss-Seidel iterations for the all-constraint + fallback path. The direct bilateral block path is controlled by + `block_iterations` and `contact_iterations`. Must be greater than zero. + Defaults to `100`. + """ + + tolerance: float = 1e-5 + """ + The convergence tolerance on the projected update size. + Must be non-negative. Defaults to `1e-5`. + """ + + regularization: float = 1e-6 + """ + Diagonal regularization added to each projected update denominator. + Must be positive. Defaults to `1e-6`. + """ + + omega: float = 1.0 + """ + Relaxation factor applied to projected Gauss-Seidel updates. + Must be in the range `(0, 2]`. Defaults to `1.0`. + """ + + block_iterations: int = 32 + """ + Number of outer DVI block iterations alternating direct bilateral solves + with projected inequality solves. Must be greater than zero. Defaults to `32`. + """ + + contact_iterations: int = 4 + """ + Number of projected Gauss-Seidel sweeps used for unilateral inequalities + during each DVI block iteration. Contacts use graph-colored sweeps on CUDA. + Must be greater than zero. Defaults to `4`. + """ + + bilateral_solve_period: int = 1 + """ + Number of DVI block iterations between repeated direct bilateral solves. + A value of `1` re-solves after every projected inequality block, preserving + the standard direct-block schedule. Must be greater than zero. Defaults to `1`. + """ + + bilateral_solver_type: Literal["LLTB", "LLTBRCM"] = "LLTB" + """ + Direct linear solver used for the bilateral constraint block. + ``LLTBRCM`` can accelerate large sparse articulated systems, while + ``LLTB`` remains preferable for small or dense systems. Defaults to + ``LLTB``. + """ + + bilateral_solver_kwargs: dict[str, Any] = field(default_factory=dict) + """ + Additional keyword arguments passed to the bilateral linear solver. + Defaults to an empty dictionary. + """ + + contact_jacobi_omega: float = 0.3 + """ + Step size for contact Jacobi updates and block-preconditioned contact + updates. Must be in the range `(0, 2]`. Defaults to `0.3`. + """ + + contact_jacobi_relaxation: float = 0.9 + """ + Solution mixing factor for contact Jacobi updates and block-preconditioned + contact updates. Must be in the range `(0, 1]`. Defaults to `0.9`. + """ + + contact_block_preconditioner: bool = False + """ + Whether to use a full 3x3 contact diagonal block preconditioner for DVI + projected contact updates. Defaults to `False`. + """ + + warmstart_mode: Literal["none", "internal", "containers"] = "containers" + """ + Warmstart mode to be used for the DVI solver. + Uses the same choices as the other dual dynamics solvers. Defaults to `containers`. + """ + + contact_warmstart_method: Literal[ + "key_and_position", + "geom_pair_net_force", + "key_and_position_with_net_force_backup", + ] = "key_and_position_with_net_force_backup" + """ + The contact warmstart method used when `warmstart_mode` is `containers`. + See :class:`WarmstarterContacts.Method` for available options. + Defaults to `key_and_position_with_net_force_backup`. + """ + + @override + @staticmethod + def register_custom_attributes(builder: ModelBuilder) -> None: + """Register DVI custom attributes supported by the Kamino USD schema. + + DVI-specific tuning options are currently Python-only. The shared + ``max_solver_iterations`` attribute is registered by + :class:`PADMMSolverConfig` and parsed by both dynamics solvers. + """ + + @override + @staticmethod + def from_model(model: Model, **kwargs: dict[str, Any]) -> DVISolverConfig: + """Creates a :class:`DVISolverConfig` from model attributes if available. + + Args: + model: The Newton model from which to parse configurations. + """ + cfg = DVISolverConfig(**kwargs) + kamino_attrs = getattr(model, "kamino", None) + if kamino_attrs is not None and hasattr(kamino_attrs, "max_solver_iterations"): + max_iterations = int(kamino_attrs.max_solver_iterations.numpy()[0]) + if max_iterations >= 0: + cfg.max_iterations = max_iterations + cfg.validate() + return cfg + + @override + def validate(self) -> None: + """Validates the current values held by this config instance.""" + from ._src.solvers.common import WarmStartMode # noqa: PLC0415 + from ._src.solvers.warmstart import WarmstarterContacts # noqa: PLC0415 + + if self.max_iterations <= 0: + raise ValueError(f"Invalid maximum iterations: {self.max_iterations}. Must be a positive integer.") + if self.tolerance < 0.0: + raise ValueError(f"Invalid tolerance: {self.tolerance}. Must be non-negative.") + if self.regularization <= 0.0: + raise ValueError(f"Invalid regularization: {self.regularization}. Must be greater than zero.") + if self.omega <= 0.0 or self.omega > 2.0: + raise ValueError(f"Invalid omega: {self.omega}. Must be in the range (0, 2].") + if self.block_iterations <= 0: + raise ValueError(f"Invalid block iterations: {self.block_iterations}. Must be a positive integer.") + if self.contact_iterations <= 0: + raise ValueError(f"Invalid contact iterations: {self.contact_iterations}. Must be a positive integer.") + if self.bilateral_solve_period <= 0: + raise ValueError( + f"Invalid bilateral solve period: {self.bilateral_solve_period}. Must be a positive integer." + ) + if self.bilateral_solver_type not in {"LLTB", "LLTBRCM"}: + raise ValueError( + f"Invalid bilateral solver type: {self.bilateral_solver_type}. Must be one of ['LLTB', 'LLTBRCM']." + ) + if self.contact_jacobi_omega <= 0.0 or self.contact_jacobi_omega > 2.0: + raise ValueError(f"Invalid contact Jacobi omega: {self.contact_jacobi_omega}. Must be in the range (0, 2].") + if self.contact_jacobi_relaxation <= 0.0 or self.contact_jacobi_relaxation > 1.0: + raise ValueError( + f"Invalid contact Jacobi relaxation: {self.contact_jacobi_relaxation}. Must be in the range (0, 1]." + ) + WarmStartMode.from_string(self.warmstart_mode) + WarmstarterContacts.Method.from_string(self.contact_warmstart_method) + implemented_contact_warmstart_methods = { + "key_and_position", + "geom_pair_net_force", + "key_and_position_with_net_force_backup", + } + if self.contact_warmstart_method not in implemented_contact_warmstart_methods: + raise ValueError( + f"DVI contact warmstart method is not implemented: {self.contact_warmstart_method}. " + f"Choose one of {sorted(implemented_contact_warmstart_methods)}." + ) + + @override + def __post_init__(self): + """Post-initialization to validate configurations.""" + self.validate() + + @dataclass class ForwardKinematicsSolverConfig: """ diff --git a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_all_geoms.py b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_all_geoms.py index 676161ae11..e025886dd9 100644 --- a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_all_geoms.py +++ b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_all_geoms.py @@ -132,7 +132,7 @@ def __init__(self, viewer: newton.viewer.ViewerBase, args=None): def capture(self): self.graph = None - if self.device.is_cuda: + if self.device.is_cuda and not wp.config.verify_cuda: with wp.ScopedCapture() as capture: self.simulate() self.graph = capture.graph diff --git a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_all_joints.py b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_all_joints.py index 22d9922c08..b153918559 100644 --- a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_all_joints.py +++ b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_all_joints.py @@ -85,7 +85,7 @@ def __init__(self, viewer: newton.viewer.ViewerBase, args=None): def capture(self): self.graph = None - if self.device.is_cuda: + if self.device.is_cuda and not wp.config.verify_cuda: with wp.ScopedCapture() as capture: self.simulate() self.graph = capture.graph diff --git a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_box_on_plane.py b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_box_on_plane.py index 9086121772..47d80d523a 100644 --- a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_box_on_plane.py +++ b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_box_on_plane.py @@ -121,7 +121,7 @@ def __init__(self, viewer: newton.viewer.ViewerBase, args=None): def capture(self): self.graph = None - if self.device.is_cuda: + if self.device.is_cuda and not wp.config.verify_cuda: with wp.ScopedCapture() as capture: self.simulate() self.graph = capture.graph diff --git a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_box_pendulum.py b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_box_pendulum.py index e0a23b5a84..799501cece 100644 --- a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_box_pendulum.py +++ b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_box_pendulum.py @@ -109,7 +109,7 @@ def __init__(self, viewer: newton.viewer.ViewerBase, args=None): def capture(self): self.graph = None - if self.device.is_cuda: + if self.device.is_cuda and not wp.config.verify_cuda: with wp.ScopedCapture() as capture: self.simulate() self.graph = capture.graph diff --git a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_boxes_hinged.py b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_boxes_hinged.py index c8a049ffc4..d28382e344 100644 --- a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_boxes_hinged.py +++ b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_boxes_hinged.py @@ -109,7 +109,7 @@ def __init__(self, viewer: newton.viewer.ViewerBase, args=None): def capture(self): self.graph = None - if self.device.is_cuda: + if self.device.is_cuda and not wp.config.verify_cuda: with wp.ScopedCapture() as capture: self.simulate() self.graph = capture.graph diff --git a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_boxes_nunchaku.py b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_boxes_nunchaku.py index 3ead9a372d..210b019492 100644 --- a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_boxes_nunchaku.py +++ b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_boxes_nunchaku.py @@ -109,7 +109,7 @@ def __init__(self, viewer: newton.viewer.ViewerBase, args=None): def capture(self): self.graph = None - if self.device.is_cuda: + if self.device.is_cuda and not wp.config.verify_cuda: with wp.ScopedCapture() as capture: self.simulate() self.graph = capture.graph diff --git a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_cartpole.py b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_cartpole.py index 944cb49460..05e29ff708 100644 --- a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_cartpole.py +++ b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_cartpole.py @@ -109,7 +109,7 @@ def __init__(self, viewer: newton.viewer.ViewerBase, args=None): def capture(self): self.graph = None - if self.device.is_cuda: + if self.device.is_cuda and not wp.config.verify_cuda: with wp.ScopedCapture() as capture: self.simulate() self.graph = capture.graph diff --git a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_sphere.py b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_sphere.py index aca1bc98cc..ab0ad69416 100644 --- a/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_sphere.py +++ b/newton/_src/solvers/kamino/examples/newton/example_kamino_basic_sphere.py @@ -109,7 +109,7 @@ def __init__(self, viewer: newton.viewer.ViewerBase, args=None): def capture(self): self.graph = None - if self.device.is_cuda: + if self.device.is_cuda and not wp.config.verify_cuda: with wp.ScopedCapture() as capture: self.simulate() self.graph = capture.graph diff --git a/newton/_src/solvers/kamino/examples/reset/example_reset_dr_legs.py b/newton/_src/solvers/kamino/examples/reset/example_reset_dr_legs.py index 644a51091d..47f7906892 100644 --- a/newton/_src/solvers/kamino/examples/reset/example_reset_dr_legs.py +++ b/newton/_src/solvers/kamino/examples/reset/example_reset_dr_legs.py @@ -136,7 +136,7 @@ def __init__( # Set gravity for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = False + self.builder.set_gravity(wp.vec3f(0.0), w) # Set solver config config = Simulator.Config() @@ -477,7 +477,7 @@ def plot(self, path: str | None = None, keep_frames: bool = False): device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/rl/test_multi_env_dr_legs.py b/newton/_src/solvers/kamino/examples/rl/test_multi_env_dr_legs.py index 60ac0dab04..c86e86a327 100644 --- a/newton/_src/solvers/kamino/examples/rl/test_multi_env_dr_legs.py +++ b/newton/_src/solvers/kamino/examples/rl/test_multi_env_dr_legs.py @@ -63,8 +63,6 @@ def run_test(num_worlds: int, num_steps: int, device): builder.max_contacts_per_pair = 8 # Cap contact budget to avoid Warp tile API shared memory bug offset = wp.transformf(0.0, 0.0, 0.265, 0.0, 0.0, 0.0, 1.0) set_uniform_body_pose_offset(builder=builder, offset=offset) - for w in range(builder.num_worlds): - builder.gravity[w].enabled = True # Create simulator msg.info("Creating simulator...") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_all_heterogeneous.py b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_all_heterogeneous.py index eb9ed5f8d5..5f27f1234b 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_all_heterogeneous.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_all_heterogeneous.py @@ -108,8 +108,9 @@ def __init__( self.builder: ModelBuilderKamino = basics.make_basics_heterogeneous_builder(ground=ground) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Set solver config config = Simulator.Config() @@ -301,7 +302,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_box_on_plane.py b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_box_on_plane.py index e9e1784c84..56c041d6ff 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_box_on_plane.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_box_on_plane.py @@ -59,7 +59,7 @@ def _control_callback( # Apply a time-dependent external force if t > t_start and t < t_end and wnc > 0: m = wp.float32(1.0) # Mass of the box - g = wp.float32(9.8067) # Gravitational acceleration + g = wp.float32(9.81) # Gravitational acceleration mu = wp.float32(0.9) # Friction coefficient f_ext = 1.1 * m * g * mu # Magnitude of the external force state_w_i_e[bid] = wp.spatial_vectorf(f_ext, 0.0, 0.0, 0.0, 0.0, 0.0) @@ -140,8 +140,9 @@ def __init__( ) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Set solver config config = Simulator.Config() @@ -342,7 +343,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_box_pendulum.py b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_box_pendulum.py index f7e46c4739..ecf7039ac8 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_box_pendulum.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_box_pendulum.py @@ -70,8 +70,9 @@ def __init__( ) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Set solver config config = Simulator.Config() @@ -295,7 +296,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_fourbar.py b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_fourbar.py index e3bc45436a..39cde0df6f 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_fourbar.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_fourbar.py @@ -226,8 +226,9 @@ def __init__( set_uniform_body_pose_offset(builder=self.builder, offset=offset) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Set solver config config = Simulator.Config() @@ -444,7 +445,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_hinged.py b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_hinged.py index aa05f25eb8..4db8554b42 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_hinged.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_hinged.py @@ -127,8 +127,9 @@ def __init__( ) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Set solver config config = Simulator.Config() @@ -329,7 +330,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_cartpole.py b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_cartpole.py index 729fc00352..6d7b135902 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_basics_cartpole.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_basics_cartpole.py @@ -150,8 +150,9 @@ def __init__( ) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Demo of printing builder contents in debug logging mode msg.info("self.builder.gravity:\n%s", self.builder.gravity) @@ -417,7 +418,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_dr_legs.py b/newton/_src/solvers/kamino/examples/sim/example_sim_dr_legs.py index 6b0cc97831..1e492898d3 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_dr_legs.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_dr_legs.py @@ -22,6 +22,7 @@ from newton._src.solvers.kamino._src.utils.io.usd import USDImporter from newton._src.solvers.kamino._src.utils.sim import SimulationLogger, Simulator, ViewerKamino from newton._src.solvers.kamino.examples import get_examples_output_path, run_headless +from newton._src.solvers.kamino.solver_kamino import SolverKamino ### # Module configs @@ -149,7 +150,8 @@ def __init__( gravity: bool = True, ground: bool = True, logging: bool = False, - linear_solver: str = "LLTB", + dynamics_solver: str = "padmm", + linear_solver: str | None = None, linear_solver_maxiter: int = 0, use_graph_conditionals: bool = False, headless: bool = False, @@ -159,9 +161,13 @@ def __init__( # Initialize target frames per second and corresponding time-steps self.fps = 50 self.frame_dt = 1.0 / self.fps - target_sim_dt = 0.01 if implicit_pd else 0.001 + target_sim_dt = self.frame_dt / 12 if dynamics_solver == "dvi" else 0.01 if implicit_pd else 0.001 self.sim_substeps = max(1, round(self.frame_dt / target_sim_dt)) self.sim_dt = self.frame_dt / self.sim_substeps + # DVI benefits from early contact detection because it solves inequality + # constraints slightly less accurately than PADMM. Contact forces remain + # zero until the shapes overlap. + dvi_contact_margin = 5.0e-4 if dynamics_solver == "dvi" else 0.0 msg.info(f"Using sim_dt = {self.sim_dt} ({self.sim_substeps} substeps per frame)") self.max_steps = max_steps @@ -198,10 +204,14 @@ def __init__( if ground: for w in range(num_worlds): add_ground_box(self.builder, world_index=w) + if dvi_contact_margin > 0.0: + for geom in self.builder.all_geoms: + geom.margin = max(geom.margin, dvi_contact_margin) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Set joint armatures, and verify that correct gains were loaded from the USD file for joint in self.builder.all_joints: @@ -211,15 +221,21 @@ def __init__( assert abs(joint.k_p_j[0] - 50.0) < 1e-4 assert abs(joint.k_d_j[0] - 1.0) < 1e-4 + if linear_solver is None: + linear_solver = "CR" if dynamics_solver == "dvi" else "LLTB" + if dynamics_solver == "dvi" and linear_solver == "CR" and linear_solver_maxiter == 0: + linear_solver_maxiter = 9 + # Parse the linear solver max iterations for iterative solvers from the command-line arguments linear_solver_kwargs = {"maxiter": linear_solver_maxiter} if linear_solver_maxiter > 0 else {} # Set solver config config = Simulator.Config() + config.solver = SolverKamino.Config(dynamics_solver=dynamics_solver) config.dt = self.sim_dt config.collision_detector.pipeline = "unified" # Select from {"primitive", "unified"} - config.solver.sparse_jacobian = False - config.solver.sparse_dynamics = False + config.solver.sparse_jacobian = dynamics_solver == "dvi" + config.solver.sparse_dynamics = dynamics_solver == "dvi" config.solver.integrator = "moreau" # Select from {"euler", "moreau"} config.solver.constraints.alpha = 0.1 config.solver.constraints.beta = 0.011 @@ -239,6 +255,20 @@ def __init__( config.solver.compute_solution_metrics = logging and not use_cuda_graph config.solver.dynamics.linear_solver_type = linear_solver config.solver.dynamics.linear_solver_kwargs = linear_solver_kwargs + config.solver.dynamics.preconditioning = dynamics_solver != "dvi" + if dynamics_solver == "dvi": + config.solver.constraints.gamma = 0.015 + config.solver.dvi.max_iterations = 200 + config.solver.dvi.tolerance = 1e-4 + config.solver.dvi.regularization = 1e-5 + config.solver.dvi.omega = 0.3 + config.solver.dvi.block_iterations = 4 + config.solver.dvi.contact_iterations = 2 + config.solver.dvi.bilateral_solve_period = 1 + config.solver.dvi.contact_jacobi_omega = 0.45 + config.solver.dvi.contact_jacobi_relaxation = 0.9 + config.solver.dvi.warmstart_mode = "containers" + config.solver.dvi.contact_warmstart_method = "key_and_position_with_net_force_backup" config.solver.padmm.use_graph_conditionals = use_graph_conditionals config.solver.angular_velocity_damping = 0.0 @@ -470,12 +500,18 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = default=None, help="Enable frame recording: 'sync' for synchronous, 'async' for asynchronous (non-blocking)", ) + parser.add_argument( + "--dynamics-solver", + default="padmm", + choices=["padmm", "dvi"], + help="Dynamics solver to use", + ) parser.add_argument( "--linear-solver", - default="LLTB", + default=None, choices=LinearSolverShorthand.values(), type=str.upper, - help="Linear solver to use", + help="Linear solver to use; defaults to LLTB for PADMM and CR for DVI", ) parser.add_argument( "--linear-solver-maxiter", default=0, type=int, help="Max number of iterations for iterative linear solvers" @@ -508,7 +544,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") @@ -519,6 +555,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device=device, use_cuda_graph=use_cuda_graph, num_worlds=args.num_worlds, + dynamics_solver=args.dynamics_solver, linear_solver=args.linear_solver, linear_solver_maxiter=args.linear_solver_maxiter, use_graph_conditionals=args.use_graph_conditionals, diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_dr_testmech.py b/newton/_src/solvers/kamino/examples/sim/example_sim_dr_testmech.py index 018a349006..05202ff29d 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_dr_testmech.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_dr_testmech.py @@ -61,8 +61,9 @@ def __init__( msg.info("total diag inertia: %f", self.builder.worlds[0].inertia_total) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Set solver config config = Simulator.Config() @@ -253,7 +254,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_test_all_geoms.py b/newton/_src/solvers/kamino/examples/sim/example_sim_test_all_geoms.py index ff56aa1458..ab84b5bc1e 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_test_all_geoms.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_test_all_geoms.py @@ -285,7 +285,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/examples/sim/example_sim_test_all_joints.py b/newton/_src/solvers/kamino/examples/sim/example_sim_test_all_joints.py index 81f818fa1e..d2fd5a694d 100644 --- a/newton/_src/solvers/kamino/examples/sim/example_sim_test_all_joints.py +++ b/newton/_src/solvers/kamino/examples/sim/example_sim_test_all_joints.py @@ -59,8 +59,9 @@ def __init__( ) # Set gravity - for w in range(self.builder.num_worlds): - self.builder.gravity[w].enabled = gravity + if not gravity: + for w in range(self.builder.num_worlds): + self.builder.set_gravity(wp.vec3f(0.0), w) # Set solver config config = Simulator.Config() @@ -254,7 +255,7 @@ def plot(self, path: str | None = None, show: bool = False, keep_frames: bool = device = wp.get_preferred_device() # Determine if CUDA graphs should be used for execution - can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) + can_use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device) and not wp.config.verify_cuda use_cuda_graph = can_use_cuda_graph and args.cuda_graph msg.info(f"can_use_cuda_graph: {can_use_cuda_graph}") msg.info(f"use_cuda_graph: {use_cuda_graph}") diff --git a/newton/_src/solvers/kamino/solver_kamino.py b/newton/_src/solvers/kamino/solver_kamino.py index 2a8bed79f6..67acbeb7fe 100644 --- a/newton/_src/solvers/kamino/solver_kamino.py +++ b/newton/_src/solvers/kamino/solver_kamino.py @@ -19,7 +19,6 @@ from ...sim import ( Contacts, Control, - JointTargetMode, JointType, Model, ModelBuilder, @@ -27,6 +26,7 @@ State, StateFlags, ) +from ...sim.collide import _estimate_rigid_contact_max from ..coupled.interface import CouplingInterface from ..solver import SolverBase @@ -36,6 +36,7 @@ ConfigBase, ConstrainedDynamicsConfig, ConstraintStabilizationConfig, + DVISolverConfig, ForwardKinematicsSolverConfig, MaterialManagerConfig, PADMMSolverConfig, @@ -58,9 +59,11 @@ class SolverKamino(SolverBase, CouplingInterface): A physics solver for simulating constrained multi-body systems containing kinematic loops, under-/overactuation, joint-limits, hard frictional contacts and restitutive impacts. - This solver uses the Proximal-ADMM algorithm to solve the forward dynamics formulated - as a Nonlinear Complementarity Problem (NCP) over the set of bilateral kinematic joint - constraints and unilateral constraints that include joint-limits and contacts. + Forward dynamics are formulated as a Nonlinear Complementarity Problem (NCP) + over bilateral kinematic joint constraints and unilateral joint-limit and + contact constraints. The default PADMM backend solves this problem with + Proximal ADMM. An opt-in DVI backend uses projected iterations with a direct + bilateral block solve. This solver is currently in Beta. @@ -106,9 +109,10 @@ class Config: A container to hold all configurations of the :class:`SolverKamino` solver. """ - sparse_jacobian: bool = False + sparse_jacobian: bool | None = None """ - Flag to indicate whether the solver should use sparse data representations for the Jacobian. + Whether to use a sparse Jacobian representation. When unspecified, defaults to `True` for DVI and `False` + for PADMM. """ sparse_dynamics: bool = False @@ -157,11 +161,18 @@ class Config: padmm: PADMMSolverConfig | None = None """ - Configurations for the dynamics solver.\n + Configurations for the PADMM dynamics solver.\n See :class:`PADMMSolverConfig` for more details.\n If `None`, default values will be used. """ + dvi: DVISolverConfig | None = None + """ + Configurations for the DVI dynamics solver.\n + See :class:`DVISolverConfig` for more details.\n + If `None`, default values will be used. + """ + fk: ForwardKinematicsSolverConfig | None = None """ Configurations for the forward kinematics solver.\n @@ -190,6 +201,13 @@ class Config: Defaults to `"euler"`. """ + dynamics_solver: Literal["padmm", "dvi"] = "padmm" + """ + The forward dynamics solver to use. Construct the config with this value + so solver-dependent defaults are initialized consistently. Defaults to + `"padmm"`. + """ + angular_velocity_damping: float = 0.0 """ A damping factor applied to the angular velocity of bodies during state integration.\n @@ -233,6 +251,7 @@ def register_custom_attributes(builder: ModelBuilder) -> None: config.ConstrainedDynamicsConfig.register_custom_attributes(builder) config.CollisionDetectorConfig.register_custom_attributes(builder) config.PADMMSolverConfig.register_custom_attributes(builder) + config.DVISolverConfig.register_custom_attributes(builder) config.MaterialManagerConfig.register_custom_attributes(builder) # Register KaminoSceneAPI custom attributes for each individual solver-level configurations @@ -286,14 +305,25 @@ def from_model(model: Model, **kwargs: dict[str, Any]) -> SolverKamino.Config: "constraints": config.ConstraintStabilizationConfig, "dynamics": config.ConstrainedDynamicsConfig, "padmm": config.PADMMSolverConfig, + "dvi": config.DVISolverConfig, "fk": config.ForwardKinematicsSolverConfig, "materials": config.MaterialManagerConfig, } for attr_name, config_cls in subconfigs.items(): nested_config = kwargs.get(attr_name, None) - nested_kwargs = nested_config.__dict__ if nested_config is not None else {} + if nested_config is not None: + nested_kwargs = nested_config.__dict__ + elif cfg.dynamics_solver == "dvi" and attr_name in {"dynamics", "dvi"}: + nested_kwargs = getattr(cfg, attr_name).__dict__ + else: + nested_kwargs = {} setattr(cfg, attr_name, config_cls.from_model(model, **nested_kwargs)) + if cfg.dynamics_solver == "dvi" and "dynamics" not in kwargs: + cfg.dynamics.preconditioning = False + + cfg.validate() + # Return the fully constructed config with sub-configurations # parsed from the model's custom attributes if available, # otherwise using defaults or provided kwargs. @@ -321,6 +351,8 @@ def validate(self) -> None: raise ValueError("Constrained dynamics config cannot be None.") elif self.padmm is None: raise ValueError("PADMM solver config cannot be None.") + elif self.dvi is None: + raise ValueError("DVI solver config cannot be None.") # Validate specialized sub-configurations # using their own built-in validations @@ -331,8 +363,20 @@ def validate(self) -> None: self.constraints.validate() self.dynamics.validate() self.padmm.validate() + self.dvi.validate() self.materials.validate() + supported_dynamics_solvers = {"padmm", "dvi"} + if self.dynamics_solver not in supported_dynamics_solvers: + raise ValueError( + f"Invalid dynamics solver: {self.dynamics_solver}. Must be one of {supported_dynamics_solvers}." + ) + if self.dynamics_solver == "dvi" and self.dynamics.preconditioning: + raise ValueError( + "The DVI solver currently requires `dynamics.preconditioning=False` so convergence checks and " + "contact cone updates stay in physical constraint units." + ) + # Conversion to JointCorrectionMode will raise an error if the input string is invalid. JointCorrectionMode.from_string(self.rotation_correction) @@ -356,6 +400,9 @@ def __post_init__(self): # Import here to avoid module-level imports and circular dependencies from . import config # noqa: PLC0415 + if self.sparse_jacobian is None: + self.sparse_jacobian = self.dynamics_solver == "dvi" + # Default-initialize any sub-configurations that were not explicitly provided by the user if self.collision_detector is None and self.use_collision_detector: self.collision_detector = config.CollisionDetectorConfig() @@ -364,9 +411,33 @@ def __post_init__(self): if self.constraints is None: self.constraints = config.ConstraintStabilizationConfig() if self.dynamics is None: - self.dynamics = config.ConstrainedDynamicsConfig() + if self.dynamics_solver == "dvi" and self.sparse_dynamics: + self.dynamics = config.ConstrainedDynamicsConfig( + preconditioning=False, + linear_solver_type="CR", + linear_solver_kwargs={"maxiter": 9}, + ) + elif self.dynamics_solver == "dvi": + self.dynamics = config.ConstrainedDynamicsConfig( + preconditioning=False, + linear_solver_type="LLTBRCM", + ) + else: + self.dynamics = config.ConstrainedDynamicsConfig() if self.padmm is None: self.padmm = config.PADMMSolverConfig() + if self.dvi is None: + if self.dynamics_solver == "dvi" and self.sparse_dynamics: + self.dvi = config.DVISolverConfig( + omega=0.3, + block_iterations=16, + contact_iterations=2, + bilateral_solve_period=2, + contact_jacobi_omega=0.45, + contact_jacobi_relaxation=0.9, + ) + else: + self.dvi = config.DVISolverConfig() if self.materials is None: self.materials = config.MaterialManagerConfig() @@ -396,7 +467,7 @@ class ResetConfig: reset_config = newton.solvers.SolverKamino.ResetConfig.to_default() solver.reset(state=state, config=reset_config) - # Preserve the current body/joint state, while resetting time, forces/torques and solver internals + # Preserve the current body state, while resetting time, forces/torques and solver internals reset_config = newton.solvers.SolverKamino.ResetConfig.preserve() solver.reset(state=state, config=reset_config) @@ -423,7 +494,7 @@ class ToDefault: @dataclass(frozen=True) class Preserve: - """Reset option, to preserve current values, assuming without check that they are consistent.""" + """Reset option, to preserve current body/base values, assuming without check that they are consistent.""" @dataclass(frozen=True) class FromJointQ: @@ -620,7 +691,23 @@ def __init__( # Store for which joints the limits are finite. This is used to validate that finiteness of limits is not changed at runtime. q_min = self._model_kamino.joints.q_j_min.numpy() q_max = self._model_kamino.joints.q_j_max.numpy() - self._built_limit_finite = (q_min > self._kamino.JOINT_QMIN) | (q_max < self._kamino.JOINT_QMAX) + built_limit_finite_np = (q_min > self._kamino.JOINT_QMIN) | (q_max < self._kamino.JOINT_QMAX) + self._built_limit_finite = wp.array( + built_limit_finite_np.astype(np.int32), + dtype=wp.int32, + device=model.device, + ) + + # Scratch array for notify validation + self._notify_violations = wp.empty(4, dtype=wp.int32, device=model.device) + + # Cache one representative shape per material. + self._material_first_shape = self._kamino.compute_material_first_shape( + self._model_kamino.geoms.material, + self._model_kamino.materials.num_materials, + ) + # Scratch scalar for material update validation + self._material_update_conflict = wp.empty(1, dtype=wp.int32, device=model.device) # Create a collision detector if enabled in the config, otherwise # set to `None` to disable internal collision detection in Kamino @@ -636,13 +723,15 @@ def __init__( if self._collision_detector_kamino is not None: self._contacts_kamino = self._collision_detector_kamino.contacts else: - # If collision detector is disabled allocate contacts manually - # TODO: We need to fix this logic to properly handle the case where the collision - # detector is disabled but contacts are still provided by Newton's collision pipeline. + # If collision detector is disabled allocate contacts based on the capacity estimate from the Newton CollisionPipeline. + world_count = self.model.world_count if self.model.rigid_contact_max == 0: - world_max_contacts = self._model_kamino.geoms.world_minimum_contacts - else: - world_max_contacts = [model.rigid_contact_max // self.model.world_count] * self.model.world_count + estimated_contacts = _estimate_rigid_contact_max(model) + # Write back to the model to ensure the CollisionPipeline capacity is consistent. + model.rigid_contact_max = ((estimated_contacts + world_count - 1) // world_count) * world_count + + # Round up to the nearest multiple of the world count to account for Kamino's per world capacity. + world_max_contacts = [(model.rigid_contact_max + world_count - 1) // world_count] * world_count self._contacts_kamino = self._kamino.ContactsKamino( # TODO: model=self._model_kamino, capacity=world_max_contacts, @@ -689,6 +778,11 @@ def reset( All state components are reset consistently with the new body poses and velocities (unless prescribed otherwise by state flags), and solver-internal buffers are cleared. + More specifically, joint coordinates and velocities are re-derived from the + resulting body state for consistency, and joint constraint forces are reset to + zero. If flags exclude :attr:`~newton.StateFlags.JOINT_Q` or + :attr:`~newton.StateFlags.JOINT_QD`, the corresponding joint coordinates or + velocities are restored after the reset instead. Args: state: The simulation state to reset (modified in place). @@ -720,22 +814,16 @@ def reset( self._model_kamino.size, self.model, state, convert_to_com_frame=False ) - # Convert body poses from origin to CoM if needed + # Convert Newton origin-frame body poses to Kamino CoM frame before reset. has_callbacks = self._solver_kamino._pre_reset_cb is not None or self._solver_kamino._post_reset_cb is not None - need_CoM_conversion = ( - not isinstance(config.body_poses, SolverKamino.ResetConfig.Preserve) - or not isinstance(config.base_pose, SolverKamino.ResetConfig.Preserve) - or has_callbacks + self._kamino.convert_body_origin_to_com( + body_com=self._model_kamino.bodies.i_r_com_i, + body_q_com=state_kamino.q_i, + body_q=state_kamino.q_i, + world_mask=world_mask if not has_callbacks else None, + body_wid=self._model_kamino.bodies.wid, ) - if need_CoM_conversion: - self._kamino.convert_body_origin_to_com( - body_com=self._model_kamino.bodies.i_r_com_i, - body_q_com=state_kamino.q_i, - body_q=state_kamino.q_i, - world_mask=world_mask if not has_callbacks else None, - body_wid=self._model_kamino.bodies.wid, - ) - # Note: we convert all worlds if callbacks are set, so they see the full state correctly + # Note: we convert all worlds if callbacks are set, so they see the full state correctly # Convert base pose from origin to CoM if needed if isinstance(config.base_pose, SolverKamino.ResetConfig.FromBaseQ): @@ -777,14 +865,13 @@ def _preserve_if_unset(array: wp.array[Any] | None, flag: int) -> None: wp.copy(array, snapshot) # Convert back body poses from COM-frame (Kamino) to body-origin frame (Newton) - if need_CoM_conversion: - self._kamino.convert_body_com_to_origin( - body_com=self._model_kamino.bodies.i_r_com_i, - body_q_com=state_kamino.q_i, - body_q=state_kamino.q_i, - world_mask=world_mask if not has_callbacks else None, - body_wid=self._model_kamino.bodies.wid, - ) + self._kamino.convert_body_com_to_origin( + body_com=self._model_kamino.bodies.i_r_com_i, + body_q_com=state_kamino.q_i, + body_q=state_kamino.q_i, + world_mask=world_mask if not has_callbacks else None, + body_wid=self._model_kamino.bodies.wid, + ) # Revert changes to config if isinstance(config.base_pose, SolverKamino.ResetConfig.FromBaseQ): @@ -875,13 +962,15 @@ def notify_model_changed(self, flags: ModelFlags | int) -> None: flags: Bitmask of :class:`~newton.ModelFlags` or custom ``int`` bits indicating which properties changed. """ self._validate_structural_invariants(flags) + self._solver_kamino.validate_model_changed(flags) if flags & (ModelFlags.JOINT_DOF_PROPERTIES | ModelFlags.ACTUATOR_PROPERTIES): # The documentation is unclear about which flag should trigger this update, so we update on both flags. self._update_actuation_types() if flags & ModelFlags.MODEL_PROPERTIES: - self._update_gravity() + # All model properties are aliased. + pass if flags & (ModelFlags.BODY_PROPERTIES | ModelFlags.BODY_INERTIAL_PROPERTIES): # q_i_0 is derived from both model.body_q and model.body_com. @@ -895,11 +984,10 @@ def notify_model_changed(self, flags: ModelFlags | int) -> None: # Geom offsets are derived from body_com and shape_transform. self._update_geom_offsets() - if flags & ModelFlags.SHAPE_PROPERTIES: - pass # TODO: contact materials. - - if flags & ModelFlags.JOINT_DOF_PROPERTIES: - pass + if flags & ModelFlags.SHAPE_PROPERTIES and self._collision_detector_kamino is not None: + # Kamino materials only need to be updated when using the Kamino collision detector. + # External Newton contacts read per-shape material values directly and don't use Kamino materials. + self._update_materials() if flags & (ModelFlags.CONSTRAINT_PROPERTIES | ModelFlags.TENDON_PROPERTIES): # Kamino does not support equality/mimic constraints or tendons, so we ignore these flags. @@ -907,6 +995,8 @@ def notify_model_changed(self, flags: ModelFlags | int) -> None: # No warning is emitted for compatibility with such an environment. pass + self._solver_kamino.notify_model_changed(flags) + handled = ( ModelFlags.MODEL_PROPERTIES | ModelFlags.BODY_PROPERTIES @@ -952,8 +1042,8 @@ def update_contacts(self, contacts: Contacts, state: State | None = None) -> Non if self._contacts_kamino is None or self._contacts_kamino.model_max_contacts_host == 0: return - # Ensure the output contacts containers has sufficient size to hold the contact data from Kamino - if self._contacts_kamino.model_max_contacts_host > contacts.rigid_contact_max: + # Kamino-generated contacts must fit in the Newton output buffer. + if self._detector is not None and self._contacts_kamino.model_max_contacts_host > contacts.rigid_contact_max: raise RuntimeError( f"Contacts container has insufficient capacity for Kamino contacts: " f"model_max_contacts={self._contacts_kamino.model_max_contacts_host} > " @@ -1129,40 +1219,23 @@ def _validate_structural_invariants(self, flags: ModelFlags | int) -> None: Raises: RuntimeError: If the solver must be recreated to apply the edit. """ - if flags & ModelFlags.JOINT_DOF_PROPERTIES: - self._check_dynamic_constraint_topology() - self._check_limit_capacity() - if flags & (ModelFlags.JOINT_DOF_PROPERTIES | ModelFlags.ACTUATOR_PROPERTIES): - self._check_actuation_types() - - def _reduce_dof_maximum_by_joint(self, values: np.ndarray) -> np.ndarray: - """Reduce per-DoF values to per-joint maxima, including zero-DoF joints.""" - starts = self.model.joint_qd_start.numpy() - nonempty = np.diff(starts) > 0 - maxima = np.zeros(self.model.joint_count, dtype=values.dtype) - if np.any(nonempty): - maxima[nonempty] = np.maximum.reduceat(values, starts[:-1][nonempty]) - return maxima - - def _check_dynamic_constraint_topology(self) -> None: - """Check that each joint retains its as-built dynamic status. - - Kamino marks a joint dynamic when any of its DoFs has positive armature, - damping, target stiffness, or target damping. A dynamic joint receives - one dynamic constraint per DoF, so preserving this status also preserves - its constraint count. - """ - dof_dynamic = ( - (self.model.joint_armature.numpy() > 0.0) - | (self.model.joint_damping.numpy() > 0.0) - | (self.model.joint_target_ke.numpy() > 0.0) - | (self.model.joint_target_kd.numpy() > 0.0) + check_dof = bool(flags & ModelFlags.JOINT_DOF_PROPERTIES) + check_actuation = bool(flags & (ModelFlags.JOINT_DOF_PROPERTIES | ModelFlags.ACTUATOR_PROPERTIES)) + if not check_dof and not check_actuation: + return + + sentinel = self._kamino.validate_model_joint_updates( + self.model, + self._model_kamino.joints, + self._built_limit_finite, + self._notify_violations, + check_dof=check_dof, + check_actuation=check_actuation, ) - current_dynamic = self._reduce_dof_maximum_by_joint(dof_dynamic) - built_dynamic = self._model_kamino.joints.num_dynamic_cts.numpy() > 0 - changed = np.flatnonzero(current_dynamic != built_dynamic) - if changed.size > 0: - joint = int(changed[0]) # report only first violation + dynamic_joint, limit_dof, actuation_joint, invalid_joint = self._notify_violations.numpy() + + if dynamic_joint != sentinel: + joint = int(dynamic_joint) raise RuntimeError( f"Changing dynamic constraint topology for joint {joint} " f"({self.model.joint_label[joint]!r}) is not supported; recreate SolverKamino to apply the change. " @@ -1170,61 +1243,27 @@ def _check_dynamic_constraint_topology(self) -> None: "The opposite is also true: if the values are updated to zero, while they were non-zero when creating the solver, the dynamic constraint topology also changes." ) - def _check_actuation_types(self) -> None: - """Check that each joint remains in its as-built actuation partition. + if limit_dof != sentinel: + dof = int(limit_dof) + raise RuntimeError( + f"Changing the existence of a joint limit for DoF {dof} " + f"is not supported; recreate SolverKamino to apply the change." + ) - The comparison follows Kamino's per-joint aggregation so individual DoF - changes are allowed when the joint remains actuated or remains passive. - """ - current_actuation = self._get_joint_actuation() - built_actuation = self._model_kamino.joints.act_type.numpy() - current_passive = current_actuation == self._kamino.JointActuationType.PASSIVE - built_passive = built_actuation == self._kamino.JointActuationType.PASSIVE - changed = np.flatnonzero(current_passive != built_passive) - if changed.size > 0: - joint = int(changed[0]) # report only first violation + if actuation_joint != sentinel: + joint = int(actuation_joint) raise RuntimeError( f"Changing the actuation partition for joint {joint} " f"({self.model.joint_label[joint]!r}) is not supported; recreate SolverKamino to apply the change." ) - def _get_joint_actuation(self) -> np.ndarray: - """Compute Kamino's per-joint actuation types. - - Newton stores one target mode per DoF. Kamino takes the maximum target - mode over all DoFs of a joint, then maps that value to a - ``JointActuationType``. Zero-DoF joints retain the reduction's default - target mode, ``JointTargetMode.NONE``. - """ - target_modes = self._reduce_dof_maximum_by_joint(self.model.joint_target_mode.numpy()) - return np.array( - [ - int(self._kamino.JointActuationType.from_newton(JointTargetMode(int(target_mode)))) - for target_mode in target_modes - ], - dtype=np.int32, - ) + if invalid_joint != sentinel: + joint = int(invalid_joint) + raise ValueError(f"Unsupported joint target mode for joint {joint}") def _update_actuation_types(self) -> None: """Refresh actuation modes without changing the passive/actuated layout.""" - self._model_kamino.joints.act_type.assign(self._get_joint_actuation()) - - def _check_limit_capacity(self) -> None: - """Check that each DoF retains its as-built finite-limit status.""" - current_finite = (self.model.joint_limit_lower.numpy() > self._kamino.JOINT_QMIN) | ( - self.model.joint_limit_upper.numpy() < self._kamino.JOINT_QMAX - ) - changed = np.flatnonzero(current_finite != self._built_limit_finite) - if changed.size > 0: - dof = int(changed[0]) # report only first violation - raise RuntimeError( - f"Changing the existence of a joint limit for DoF {dof} " - f"is not supported; recreate SolverKamino to apply the change." - ) - - def _update_gravity(self): - """Update Kamino's :class:`GravityModel` from Newton's ``model.gravity``.""" - self._kamino.convert_model_gravity(self.model, self._model_kamino.gravity) + self._kamino.convert_model_joint_actuation(self.model, self._model_kamino.joints) def _update_body_initial_pose(self): """Recompute Kamino's CoM-frame initial body poses.""" @@ -1246,3 +1285,12 @@ def _update_geom_offsets(self): def _update_joint_transforms(self): """Re-derive Kamino joint anchors and axes from Newton's joint transforms.""" self._kamino.convert_model_joint_transforms(self.model, self._model_kamino.joints) + + def _update_materials(self) -> None: + """Refresh Kamino contact-material tables using cached representative shapes.""" + self._kamino.convert_model_materials( + self.model, + self._model_kamino, + self._material_first_shape, + self._material_update_conflict, + ) diff --git a/newton/_src/solvers/kamino/tests/__init__.py b/newton/_src/solvers/kamino/tests/__init__.py index 5826512f87..28a13c045f 100644 --- a/newton/_src/solvers/kamino/tests/__init__.py +++ b/newton/_src/solvers/kamino/tests/__init__.py @@ -43,7 +43,7 @@ class TestContext: ### -def setup_tests(verbose: bool = False, device: wp.DeviceLike | str | None = None, clear_cache: bool = True): +def setup_tests(verbose: bool = False, device: wp.DeviceLike | str | None = None, clear_cache: bool = False): # Numpy configuration np.set_printoptions( linewidth=999999, edgeitems=999999, threshold=999999, precision=10, suppress=True diff --git a/newton/_src/solvers/kamino/tests/__main__.py b/newton/_src/solvers/kamino/tests/__main__.py index 40eb82d241..30281e984e 100644 --- a/newton/_src/solvers/kamino/tests/__main__.py +++ b/newton/_src/solvers/kamino/tests/__main__.py @@ -53,7 +53,7 @@ class ModuleHeaderTestRunner(unittest.TextTestRunner): ) parser.add_argument( "--clear-cache", - default=True, # Edit to enable/disable cache clear (if not running in command line) + default=False, # Edit to enable/disable cache clear (if not running in command line) action=argparse.BooleanOptionalAction, help="Whether to clear the warp cache before running tests.", ) diff --git a/newton/_src/solvers/kamino/tests/test_core_builder.py b/newton/_src/solvers/kamino/tests/test_core_builder.py index abe97f2e33..4cccfe44f9 100644 --- a/newton/_src/solvers/kamino/tests/test_core_builder.py +++ b/newton/_src/solvers/kamino/tests/test_core_builder.py @@ -14,11 +14,6 @@ from newton._src.solvers.kamino._src.core.bodies import RigidBodyDescriptor from newton._src.solvers.kamino._src.core.builder import ModelBuilderKamino from newton._src.solvers.kamino._src.core.geometry import GeometryDescriptor -from newton._src.solvers.kamino._src.core.gravity import ( - GRAVITY_ACCEL_DEFAULT, - GRAVITY_DIREC_DEFAULT, - GRAVITY_NAME_DEFAULT, -) from newton._src.solvers.kamino._src.core.joints import JointActuationType, JointDescriptor, JointDoFType from newton._src.solvers.kamino._src.core.materials import MaterialDescriptor from newton._src.solvers.kamino._src.core.model import ModelKamino @@ -188,6 +183,7 @@ def test_01_make_default_with_world(self): self.assertEqual(len(builder.geoms), 1) self.assertEqual(len(builder.geoms[0]), 0) self.assertEqual(len(builder.materials), 1) # Default material is always created + np.testing.assert_array_equal(builder.gravity[0].vector, np.array([0.0, 0.0, -9.81], dtype=np.float32)) def test_02_add_world(self): builder = ModelBuilderKamino() @@ -197,9 +193,32 @@ def test_02_add_world(self): self.assertEqual(builder.worlds[wid].wid, wid) self.assertEqual(builder.worlds[wid].name, "test_world") self.assertEqual(builder.up_axes[wid], Axis.Y) - self.assertEqual(builder.gravity[wid].name, GRAVITY_NAME_DEFAULT) - self.assertEqual(builder.gravity[wid].acceleration, GRAVITY_ACCEL_DEFAULT) - np.testing.assert_array_equal(builder.gravity[wid].direction, np.array(GRAVITY_DIREC_DEFAULT, dtype=np.float32)) + np.testing.assert_array_equal(builder.gravity[wid].vector, np.array([0.0, -9.81, 0.0], dtype=np.float32)) + + def test_add_world_accepts_arraylike_gravity(self): + """Store a list gravity vector when adding a world.""" + builder = ModelBuilderKamino() + + wid = builder.add_world(gravity=[1.0, -2.0, 3.0]) + + np.testing.assert_array_equal(builder.gravity[wid].vector, np.array([1.0, -2.0, 3.0], dtype=np.float32)) + + def test_set_gravity_accepts_numpy_vector(self): + """Store a NumPy gravity vector after adding a world.""" + builder = ModelBuilderKamino() + builder.add_world() + + builder.set_gravity(np.array([1.0, -2.0, 3.0], dtype=np.float32)) + + np.testing.assert_array_equal(builder.gravity[0].vector, np.array([1.0, -2.0, 3.0], dtype=np.float32)) + + def test_set_gravity_rejects_invalid_vector_shape(self): + """Reject gravity vectors without exactly three components.""" + builder = ModelBuilderKamino() + builder.add_world() + + with self.assertRaisesRegex(ValueError, r"shape \(3,\)"): + builder.set_gravity([0.0, -9.81]) def test_03_add_rigid_body(self): builder = ModelBuilderKamino() @@ -401,9 +420,7 @@ def test_11_add_material(self): wid = builder.add_world(name="test_world", up_axis=Axis.Z) self.assertEqual(builder.num_materials, 1) # Default material exists - material = MaterialDescriptor( - name="test_material", density=500.0, restitution=0.8, static_friction=0.6, dynamic_friction=0.4 - ) + material = MaterialDescriptor(name="test_material", restitution=0.8, static_friction=0.6, dynamic_friction=0.4) mid = builder.add_material(material=material) self.assertEqual(builder.num_materials, 2) @@ -411,7 +428,6 @@ def test_11_add_material(self): self.assertEqual(mid, builder.materials[mid].mid) self.assertEqual(builder.materials[mid].name, "test_material") self.assertEqual(builder.materials[mid].wid, wid) - self.assertEqual(builder.materials[mid].density, 500.0) self.assertEqual(builder.materials[mid].restitution, 0.8) self.assertEqual(builder.materials[mid].static_friction, 0.6) self.assertEqual(builder.materials[mid].dynamic_friction, 0.4) diff --git a/newton/_src/solvers/kamino/tests/test_core_model.py b/newton/_src/solvers/kamino/tests/test_core_model.py index f68747ee23..101adef167 100644 --- a/newton/_src/solvers/kamino/tests/test_core_model.py +++ b/newton/_src/solvers/kamino/tests/test_core_model.py @@ -277,8 +277,7 @@ def test_01_model_conversions_fourbar_from_usd(self): model_newton: Model = builder_newton.finalize(skip_validation_joints=True, device=self.default_device) model_kamino: ModelKamino = builder_kamino.finalize(device=self.default_device) model_kamino_converted: ModelKamino = ModelKamino.from_newton(model_newton) - excluded = ["base_joint_index"] - test_util_checks.assert_model_equal(self, model_kamino_converted, model_kamino, excluded=excluded) + test_util_checks.assert_model_equal(self, model_kamino_converted, model_kamino) # TODO: IMPLEMENT THIS CHECK: We wanna see if the both generate # the same data containers and unilateral constraint info @@ -338,8 +337,9 @@ def test_02_model_conversions_dr_testmech_from_usd(self): # so inv_i_I_i needs a somewhat higher tolerance. rtol = {"inv_i_I_i": 1e-5} atol = {"inv_i_I_i": 1e-6} + excluded = ["ptr"] test_util_checks.assert_model_equal( - self, model_kamino_converted, model_kamino, excluded=["ptr"], rtol=rtol, atol=atol + self, model_kamino_converted, model_kamino, excluded=excluded, rtol=rtol, atol=atol ) def test_03_model_conversions_dr_legs_from_usd(self): @@ -395,7 +395,7 @@ def test_03_model_conversions_dr_legs_from_usd(self): # geom-pairs of joint neighbours to `shape_collision_filter_pairs` regardless of # whether they are actually collidable or not, which leads to differences in the # number of excluded pairs and their contents - excluded = ["base_joint_index", "ptr", "group", "gap", "num_excluded_pairs", "excluded_pairs"] + excluded = ["ptr", "group", "gap", "num_excluded_pairs", "excluded_pairs"] rtol = {"inv_i_I_i": 1e-5} atol = {"inv_i_I_i": 1e-6} test_util_checks.assert_model_equal( @@ -470,6 +470,43 @@ def test_04_model_conversions_anymal_d_from_usd(self): ] test_util_checks.assert_model_equal(self, model_kamino_converted, model_kamino, excluded=excluded) + def test_05_model_conversions_base_assignment_non_floating_root(self): + """ + Test per-world base assignment when articulation roots are not unary free joints. + + A free-rooted articulation following a fixed-rooted one still provides the + world's floating base without warning; a world whose articulations are + fixed-rooted or rooted by a free joint with a body parent gets no base and + warns that floating base resets are disabled. + """ + + def build_model(free_root: str | None) -> Model: + builder: ModelBuilder = ModelBuilder() + SolverKamino.register_custom_attributes(builder) + body_fixed = builder.add_link() + builder.add_shape_box(body_fixed) + joint_fixed = builder.add_joint_fixed(parent=-1, child=body_fixed) + builder.add_articulation([joint_fixed]) + if free_root is not None: + body_free = builder.add_link(xform=wp.transform(wp.vec3(0.0, 0.0, 2.0), wp.quat_identity())) + builder.add_shape_box(body_free) + parent = -1 if free_root == "world" else body_fixed + joint_free = builder.add_joint_free(child=body_free, parent=parent) + builder.add_articulation([joint_free]) + return builder.finalize(device=self.default_device) + + with self.assertNoLogs(level="WARNING"): + model_kamino = ModelKamino.from_newton(build_model(free_root="world")) + self.assertEqual(model_kamino.info.base_body_index.numpy().tolist(), [1]) + self.assertEqual(model_kamino.info.base_joint_index.numpy().tolist(), [1]) + + for free_root in (None, "body"): + with self.assertLogs(level="WARNING") as logs: + model_kamino = ModelKamino.from_newton(build_model(free_root=free_root)) + self.assertTrue(any("not a free joint attached to the world" in message for message in logs.output)) + self.assertEqual(model_kamino.info.base_body_index.numpy().tolist(), [-1]) + self.assertEqual(model_kamino.info.base_joint_index.numpy().tolist(), [-1]) + def test_10_model_conversions_arbitrary_axis(self): """ Test that Newton→Kamino conversion succeeds for a revolute joint @@ -634,6 +671,35 @@ def test_11_model_conversions_q_i_0_com_frame(self): np.testing.assert_allclose(q_i_0_np[2, :3], [1.1, -0.3, 0.2], atol=1e-6) np.testing.assert_allclose(q_i_0_np[2, 3:7], body_q_np[2, 3:7], atol=1e-6) + def _build_single_floating_body_com_offset_model(self) -> Model: + """Build a single floating body with a non-zero COM offset.""" + builder_newton: ModelBuilder = ModelBuilder() + SolverKamino.register_custom_attributes(builder_newton) + builder_newton.default_shape_cfg.margin = 0.0 + builder_newton.default_shape_cfg.gap = 0.0 + + builder_newton.begin_world() + + bid = builder_newton.add_link( + label="body0", + mass=1.0, + xform=wp.transformf(wp.vec3f(0.0, 0.0, 1.0), wp.quat_identity(dtype=wp.float32)), + com=wp.vec3f(0.1, 0.0, 0.0), + lock_inertia=True, + ) + builder_newton.add_shape_box(label="box0", body=bid, hx=0.05, hy=0.05, hz=0.05) + builder_newton.add_joint_free( + label="world_to_body0", + parent=-1, + child=bid, + parent_xform=wp.transform_identity(dtype=wp.float32), + child_xform=wp.transform_identity(dtype=wp.float32), + ) + + builder_newton.end_world() + + return builder_newton.finalize(skip_validation_joints=True, device=self.default_device) + def _build_com_offset_model(self, with_base_joint: bool = True): """Build a 3-body chain with non-zero COM offsets for reset tests.""" builder_newton: ModelBuilder = ModelBuilder() @@ -813,6 +879,37 @@ def test_13_base_reset_produces_body_origin_frame(self): err_msg=f"Base reset (translated): body {i} rotation mismatch", ) + def test_15_preserve_reset_keeps_joint_q_consistent_with_com_offset_body(self): + """ + Verify preserve reset leaves body_q and joint_q unchanged for COM-offset bodies. + + For a single floating body with a non-zero center-of-mass offset, a preserve + reset should not modify ``body_q`` or re-derived ``joint_q``. + """ + model = self._build_single_floating_body_com_offset_model() + solver = SolverKamino(model) + + state: State = model.state() + solver.reset(state=state) + + body_q_before = state.body_q.numpy().copy() + joint_q_before = state.joint_q.numpy().copy() + + solver.reset(state=state, config=SolverKamino.ResetConfig.preserve()) + + np.testing.assert_allclose( + state.body_q.numpy(), + body_q_before, + atol=1e-6, + err_msg="Preserve reset should not modify body_q", + ) + np.testing.assert_allclose( + state.joint_q.numpy(), + joint_q_before, + atol=1e-6, + err_msg="Preserve reset should not modify joint_q when body_q is unchanged", + ) + def test_14_model_conversions_shape_offset_com_relative(self): """ Test that ``geoms.offset`` stores COM-relative shape positions diff --git a/newton/_src/solvers/kamino/tests/test_core_shapes.py b/newton/_src/solvers/kamino/tests/test_core_shapes.py index e13eeb3a28..a48ca1fdf8 100644 --- a/newton/_src/solvers/kamino/tests/test_core_shapes.py +++ b/newton/_src/solvers/kamino/tests/test_core_shapes.py @@ -21,6 +21,7 @@ MeshShape, PlaneShape, SphereShape, + _max_contacts_for_shape_pair_impl, ) from newton._src.solvers.kamino._src.utils import logger as msg from newton._src.solvers.kamino.tests import setup_tests, test_context @@ -30,6 +31,21 @@ ### +class TestMaxContactsForShapePair(unittest.TestCase): + def test_catch_canonicalization_order_errors(self): + """Verify the canonical implementation rejects every reversed shape pair. + + This catches accidentally declaring the shape pair in the wrong order. + """ + for type_a in GeoType: + for type_b in GeoType: + if type_a <= type_b: + continue + + with self.subTest(type_a=type_a, type_b=type_b): + self.assertEqual(_max_contacts_for_shape_pair_impl(int(type_a), int(type_b)), (0, 0)) + + class TestShapeDescriptors(unittest.TestCase): def setUp(self): if not test_context.setup_done: diff --git a/newton/_src/solvers/kamino/tests/test_core_world.py b/newton/_src/solvers/kamino/tests/test_core_world.py index 040f5883d4..389895c337 100644 --- a/newton/_src/solvers/kamino/tests/test_core_world.py +++ b/newton/_src/solvers/kamino/tests/test_core_world.py @@ -12,12 +12,7 @@ from newton._src.geometry.types import GeoType from newton._src.solvers.kamino._src.core.bodies import RigidBodyDescriptor from newton._src.solvers.kamino._src.core.geometry import GeometryDescriptor -from newton._src.solvers.kamino._src.core.gravity import ( - GRAVITY_ACCEL_DEFAULT, - GRAVITY_DIREC_DEFAULT, - GRAVITY_NAME_DEFAULT, - GravityDescriptor, -) +from newton._src.solvers.kamino._src.core.gravity import GravityDescriptor from newton._src.solvers.kamino._src.core.joints import ( JOINT_DQMAX, JOINT_QMAX, @@ -28,7 +23,6 @@ JointDoFType, ) from newton._src.solvers.kamino._src.core.materials import ( - DEFAULT_DENSITY, DEFAULT_FRICTION, DEFAULT_RESTITUTION, MaterialDescriptor, @@ -44,71 +38,19 @@ class TestGravityDescriptor(unittest.TestCase): - def setUp(self): - if not test_context.setup_done: - setup_tests(clear_cache=False) - self.default_device = wp.get_device(test_context.device) - self.verbose = test_context.verbose + def test_default_vector(self): + """Use Newton's default gravity along negative Z.""" + gravity = GravityDescriptor(name="gravity") - # Set debug-level logging to print verbose test output to console - if self.verbose: - print("\n") # Add newline before test output for better readability - msg.set_log_level(msg.LogLevel.DEBUG) - else: - msg.reset_log_level() + self.assertEqual(gravity.name, "gravity") + np.testing.assert_array_equal(gravity.vector, np.array([0.0, 0.0, -9.81], dtype=np.float32)) - def tearDown(self): - self.default_device = None - if self.verbose: - msg.reset_log_level() + def test_custom_vector(self): + """Store a custom gravity vector directly.""" + gravity = GravityDescriptor(vector=wp.vec3f(1.0, -2.0, 3.0), name="custom") - def test_00_default_construction(self): - gravity = GravityDescriptor() - msg.info(f"gravity: {gravity}") - self.assertIsInstance(gravity, GravityDescriptor) - self.assertEqual(gravity.name, GRAVITY_NAME_DEFAULT) - self.assertEqual(gravity.enabled, True) - self.assertEqual(gravity.acceleration, GRAVITY_ACCEL_DEFAULT) - expected_direction = np.array(GRAVITY_DIREC_DEFAULT, dtype=np.float32) - expected_dir_accel = np.array([*GRAVITY_DIREC_DEFAULT, GRAVITY_ACCEL_DEFAULT], dtype=np.float32) - expected_vector = np.array([0.0, 0.0, -GRAVITY_ACCEL_DEFAULT, 1.0], dtype=np.float32) - np.testing.assert_array_equal(gravity.direction, expected_direction) - np.testing.assert_array_equal(gravity.dir_accel(), expected_dir_accel) - np.testing.assert_array_equal(gravity.vector(), expected_vector) - - def test_01_with_parameters_and_dir_as_list(self): - gravity = GravityDescriptor(name="test_gravity", enabled=False, acceleration=15.0, direction=[1.0, 0.0, 0.0]) - msg.info(f"gravity: {gravity}") - self.assertIsInstance(gravity, GravityDescriptor) - self.assertEqual(gravity.name, "test_gravity") - self.assertEqual(gravity.enabled, False) - self.assertEqual(gravity.acceleration, 15.0) - np.testing.assert_array_equal(gravity.direction, np.array([1.0, 0.0, 0.0], dtype=np.float32)) - np.testing.assert_array_equal(gravity.dir_accel(), np.array([1.0, 0.0, 0.0, 15.0], dtype=np.float32)) - np.testing.assert_array_equal(gravity.vector(), np.array([15.0, 0.0, 0.0, 0.0], dtype=np.float32)) - - def test_02_with_parameters_and_dir_as_tuple(self): - gravity = GravityDescriptor(name="test_gravity", enabled=False, acceleration=9.0, direction=(1.0, 0.0, 0.0)) - msg.info(f"gravity: {gravity}") - self.assertIsInstance(gravity, GravityDescriptor) - self.assertEqual(gravity.name, "test_gravity") - self.assertEqual(gravity.enabled, False) - self.assertEqual(gravity.acceleration, 9.0) - np.testing.assert_array_equal(gravity.direction, np.array([1.0, 0.0, 0.0], dtype=np.float32)) - np.testing.assert_array_equal(gravity.dir_accel(), np.array([1.0, 0.0, 0.0, 9.0], dtype=np.float32)) - np.testing.assert_array_equal(gravity.vector(), np.array([9.0, 0.0, 0.0, 0.0], dtype=np.float32)) - - def test_03_with_parameters_and_dir_as_nparray(self): - direction = np.array([1.0, 0.0, 0.0], dtype=np.float32) - gravity = GravityDescriptor(name="test_gravity", enabled=False, acceleration=12.0, direction=direction) - msg.info(f"gravity: {gravity}") - self.assertIsInstance(gravity, GravityDescriptor) - self.assertEqual(gravity.name, "test_gravity") - self.assertEqual(gravity.enabled, False) - self.assertEqual(gravity.acceleration, 12.0) - np.testing.assert_array_equal(gravity.direction, np.array([1.0, 0.0, 0.0], dtype=np.float32)) - np.testing.assert_array_equal(gravity.dir_accel(), np.array([1.0, 0.0, 0.0, 12.0], dtype=np.float32)) - np.testing.assert_array_equal(gravity.vector(), np.array([12.0, 0.0, 0.0, 0.0], dtype=np.float32)) + self.assertEqual(gravity.name, "custom") + np.testing.assert_array_equal(gravity.vector, np.array([1.0, -2.0, 3.0], dtype=np.float32)) class TestBodyDescriptor(unittest.TestCase): @@ -383,7 +325,6 @@ def test_00_default_construction(self): self.assertIsInstance(mat, MaterialDescriptor) self.assertEqual(mat.name, "test_mat") - self.assertEqual(mat.density, DEFAULT_DENSITY) self.assertEqual(mat.restitution, DEFAULT_RESTITUTION) self.assertEqual(mat.static_friction, DEFAULT_FRICTION) self.assertEqual(mat.dynamic_friction, DEFAULT_FRICTION) @@ -393,7 +334,6 @@ def test_00_default_construction(self): def test_01_with_properties(self): mat = MaterialDescriptor( name="test_mat", - density=500.0, restitution=0.5, static_friction=0.6, dynamic_friction=0.4, @@ -402,7 +342,6 @@ def test_01_with_properties(self): self.assertIsInstance(mat, MaterialDescriptor) self.assertEqual(mat.name, "test_mat") - self.assertEqual(mat.density, 500.0) self.assertEqual(mat.restitution, 0.5) self.assertEqual(mat.static_friction, 0.6) self.assertEqual(mat.dynamic_friction, 0.4) diff --git a/newton/_src/solvers/kamino/tests/test_geometry_contacts.py b/newton/_src/solvers/kamino/tests/test_geometry_contacts.py index 2325e32d6a..5a52d542bc 100644 --- a/newton/_src/solvers/kamino/tests/test_geometry_contacts.py +++ b/newton/_src/solvers/kamino/tests/test_geometry_contacts.py @@ -1144,7 +1144,7 @@ def test_05_multi_world(self): self._seed_constant_linear_force(contacts, f_world) - kamino_out = ContactsKamino(capacity=nc_orig + 32, device=self.default_device) + kamino_out = ContactsKamino(capacity=[9 + 10, 9 + 10, 4 + 10], device=self.default_device) convert_contacts_newton_to_kamino(model, state, contacts, kamino_out, convert_forces=True) nc_kamino = int(kamino_out.model_active_contacts.numpy()[0]) self.assertGreater(nc_kamino, 0) @@ -1392,6 +1392,51 @@ def test_08_edge_cases(self): convert_forces=True, ) + def test_09_multi_world_per_world_capacity_not_starved(self): + """Preserve capacity for later worlds when earlier worlds are saturated. + + The conversion must examine all Newton contact slots so a saturated + world cannot prevent a later world from filling its own contact cap. + + Test the following regression: + Scanning with the lower Kamino contact capacity would miss contacts. + """ + scene = ModelBuilder() + scene.add_ground_plane() + scene.add_world(build_nunchaku_scene(ground=False)) + + single_box = ModelBuilder() + body = single_box.add_link() + no_gap = ModelBuilder.ShapeConfig(gap=0.0) + single_box.add_shape_box(body, hx=0.25, hy=0.25, hz=0.25, cfg=no_gap) + joint = single_box.add_joint_free( + parent=-1, + child=body, + parent_xform=wp.transform(p=wp.vec3(0.0, 0.0, 0.25), q=wp.quat_identity()), + child_xform=wp.transform_identity(), + ) + single_box.add_articulation([joint]) + scene.add_world(single_box, xform=wp.transform(p=wp.vec3(10.0, 0.0, 0.0))) + + model = scene.finalize(self.default_device) + self.assertEqual(model.world_count, 2) + state = model.state() + newton.eval_fk(model, model.joint_q, model.joint_qd, state) + collision_pipeline = newton.CollisionPipeline(model) + contacts = collision_pipeline.contacts() + collision_pipeline.collide(state, contacts) + + input_count = int(contacts.rigid_contact_count.numpy()[0]) + self.assertGreater(input_count, 2) + + kamino = ContactsKamino(capacity=[1, 1], device=self.default_device) + convert_contacts_newton_to_kamino(model, state, contacts, kamino) + + np.testing.assert_array_equal(kamino.world_active_contacts.numpy(), [1, 1]) + self.assertEqual(int(kamino.model_active_contacts.numpy()[0]), 2) + bid_AB = kamino.bid_AB.numpy()[:2] + self.assertTrue(np.all(bid_AB[:, 1] >= 0)) + ### # Test execution diff --git a/newton/_src/solvers/kamino/tests/test_geometry_detector.py b/newton/_src/solvers/kamino/tests/test_geometry_detector.py index c449af414f..3c5b2202f9 100644 --- a/newton/_src/solvers/kamino/tests/test_geometry_detector.py +++ b/newton/_src/solvers/kamino/tests/test_geometry_detector.py @@ -8,9 +8,15 @@ import numpy as np import warp as wp +from newton._src.solvers.kamino._src.core.builder import ModelBuilderKamino from newton._src.solvers.kamino._src.geometry import ( CollisionDetector, ) +from newton._src.solvers.kamino._src.geometry.detector import ( + _cap_world_contacts_at_total, + _estimate_fallback_world_max_contacts, + _resolve_contact_capacity, +) from newton._src.solvers.kamino._src.models.builders import basics from newton._src.solvers.kamino._src.models.builders.utils import make_homogeneous_builder from newton._src.solvers.kamino._src.utils import logger as msg @@ -161,6 +167,83 @@ def test_02_unified_pipeline(self): header="unified pipeline", ) + def test_03_pair_based_world_budgets(self): + """Verify per-world budgets follow geometry rather than an equal model split.""" + builder = make_homogeneous_builder(num_worlds=3, build_fn=self.build_func) + model = builder.finalize(self.default_device) + + config = CollisionDetector.Config(pipeline="primitive", broadphase="explicit") + detector = CollisionDetector(model=model, config=config) + + self.assertEqual(detector.world_max_contacts, model.geoms.world_minimum_contacts) + self.assertEqual(detector.model_max_contacts, sum(model.geoms.world_minimum_contacts)) + + def test_04_max_contacts_caps_model_total(self): + """Verify ``max_contacts`` caps rather than floors the geometry-based estimate.""" + builder = make_homogeneous_builder(num_worlds=3, build_fn=self.build_func) + model = builder.finalize(self.default_device) + uncapped_total = sum(model.geoms.world_minimum_contacts) + self.assertGreater(uncapped_total, 15) + + config = CollisionDetector.Config( + pipeline="primitive", + broadphase="explicit", + max_contacts=15, + ) + detector = CollisionDetector(model=model, config=config) + + self.assertEqual(detector.model_max_contacts, 15) + self.assertEqual(sum(detector.world_max_contacts), 15) + + def test_05_heterogeneous_world_budgets(self): + """Verify worlds with different geometry receive different contact budgets.""" + builder = ModelBuilderKamino(default_world=False) + basics.build_boxes_nunchaku(builder=builder) + builder.add_world(name="empty_world") + model = builder.finalize(self.default_device) + + config = CollisionDetector.Config(pipeline="primitive", broadphase="explicit") + detector = CollisionDetector(model=model, config=config) + + self.assertEqual(len(detector.world_max_contacts), 2) + self.assertGreater(detector.world_max_contacts[0], 0) + self.assertEqual(detector.world_max_contacts[1], 0) + self.assertEqual(detector.world_max_contacts, model.geoms.world_minimum_contacts) + + +class TestCollisionDetectorContactCapacity(unittest.TestCase): + def setUp(self): + if not test_context.setup_done: + setup_tests(clear_cache=False) + self.default_device = wp.get_device(test_context.device) + + def test_00_cap_world_contacts_at_total(self): + """Verify proportional capping preserves the configured model total.""" + capped = _cap_world_contacts_at_total([100, 50, 50], 120) + self.assertEqual(sum(capped), 120) + self.assertEqual(capped[0], 60) + + def test_01_fallback_explicit_per_world_pair_counts(self): + """Verify fallback explicit broad-phase estimates accumulate pairs per world.""" + builder = make_homogeneous_builder(num_worlds=2, build_fn=basics.build_boxes_nunchaku) + model = builder.finalize(self.default_device) + config = CollisionDetector.Config(broadphase="explicit") + + world_max = _estimate_fallback_world_max_contacts(model, config) + self.assertEqual(len(world_max), 2) + self.assertGreater(world_max[0], 0) + self.assertEqual(world_max[0], world_max[1]) + + def test_02_resolve_contact_capacity_uses_pair_metadata(self): + """Verify resolved budgets match pair-based metadata when available.""" + builder = basics.build_boxes_nunchaku() + model = builder.finalize(self.default_device) + config = CollisionDetector.Config(max_contacts=10_000) + + model_max, world_max = _resolve_contact_capacity(model, config) + self.assertEqual(world_max, model.geoms.world_minimum_contacts) + self.assertEqual(model_max, model.geoms.model_minimum_contacts) + ### # Test execution diff --git a/newton/_src/solvers/kamino/tests/test_geometry_primitive.py b/newton/_src/solvers/kamino/tests/test_geometry_primitive.py index 233cfa100f..1c1c984ec1 100644 --- a/newton/_src/solvers/kamino/tests/test_geometry_primitive.py +++ b/newton/_src/solvers/kamino/tests/test_geometry_primitive.py @@ -918,26 +918,26 @@ def test_06_box_on_box_eight_points(self): "bid_AB": np.tile(np.array([0, 1], dtype=np.int32), reps=(8, 1)), "position_A": np.array( [ - [-0.207107, -0.5, 0.5 * abs(distance)], + [0.5, -0.207107, 0.5 * abs(distance)], [0.207107, -0.5, 0.5 * abs(distance)], [-0.5, -0.207107, 0.5 * abs(distance)], - [-0.5, 0.207107, 0.5 * abs(distance)], + [-0.207107, -0.5, 0.5 * abs(distance)], [0.5, 0.207107, 0.5 * abs(distance)], - [0.5, -0.207107, 0.5 * abs(distance)], [0.207107, 0.5, 0.5 * abs(distance)], + [-0.5, 0.207107, 0.5 * abs(distance)], [-0.207107, 0.5, 0.5 * abs(distance)], ], dtype=np.float32, ), "position_B": np.array( [ - [-0.207107, -0.5, -0.5 * abs(distance)], + [0.5, -0.207107, -0.5 * abs(distance)], [0.207107, -0.5, -0.5 * abs(distance)], [-0.5, -0.207107, -0.5 * abs(distance)], - [-0.5, 0.207107, -0.5 * abs(distance)], + [-0.207107, -0.5, -0.5 * abs(distance)], [0.5, 0.207107, -0.5 * abs(distance)], - [0.5, -0.207107, -0.5 * abs(distance)], [0.207107, 0.5, -0.5 * abs(distance)], + [-0.5, 0.207107, -0.5 * abs(distance)], [-0.207107, 0.5, -0.5 * abs(distance)], ], dtype=np.float32, diff --git a/newton/_src/solvers/kamino/tests/test_kinematics_joints.py b/newton/_src/solvers/kamino/tests/test_kinematics_joints.py index b6f0c4359d..06705054fc 100644 --- a/newton/_src/solvers/kamino/tests/test_kinematics_joints.py +++ b/newton/_src/solvers/kamino/tests/test_kinematics_joints.py @@ -10,7 +10,7 @@ import warp as wp from newton._src.solvers.kamino._src.core.data import DataKamino -from newton._src.solvers.kamino._src.core.math import quat_exp, screw, screw_angular, screw_linear +from newton._src.solvers.kamino._src.core.math import quat_exp from newton._src.solvers.kamino._src.core.model import ModelKamino from newton._src.solvers.kamino._src.kinematics.joints import JointActuationType, compute_joints_data from newton._src.solvers.kamino._src.models.builders.testing import build_unary_revolute_joint_test @@ -74,8 +74,8 @@ def _set_joint_follower_body_state( r_B = wp.transform_get_translation(p_B) q_B = wp.transform_get_rotation(p_B) R_B = wp.quat_to_matrix(q_B) - v_B = screw_linear(u_B) - omega_B = screw_angular(u_B) + v_B = wp.spatial_top(u_B) + omega_B = wp.spatial_bottom(u_B) # Define the joint rotation offset j_dR_yz_j = wp.vec3f(0.0, THETA_Y_J, THETA_Z_J) # Joint residual as rotation vector @@ -107,7 +107,7 @@ def _set_joint_follower_body_state( # Offset the bose of the body by a fixed amount state_body_q_i[bid_F] = wp.transformation(r_F_new, q_F_new, dtype=wp.float32) - state_body_u_i[bid_F] = screw(v_F_new, omega_F_new) + state_body_u_i[bid_F] = wp.spatial_vectorf(*v_F_new, *omega_F_new) ### diff --git a/newton/_src/solvers/kamino/tests/test_kinematics_limits.py b/newton/_src/solvers/kamino/tests/test_kinematics_limits.py index c3762d466f..d9166976e5 100644 --- a/newton/_src/solvers/kamino/tests/test_kinematics_limits.py +++ b/newton/_src/solvers/kamino/tests/test_kinematics_limits.py @@ -11,7 +11,7 @@ import warp as wp from newton._src.solvers.kamino._src.core.data import DataKamino -from newton._src.solvers.kamino._src.core.math import quat_exp, screw, screw_angular, screw_linear +from newton._src.solvers.kamino._src.core.math import quat_exp from newton._src.solvers.kamino._src.core.model import ModelKamino from newton._src.solvers.kamino._src.kinematics.joints import compute_joints_data from newton._src.solvers.kamino._src.kinematics.limits import LimitsKamino @@ -75,8 +75,8 @@ def _set_joint_follower_body_state( R_B = wp.quat_to_matrix(q_B) # Extract the linear and angular velocity of the Base body - v_B = screw_linear(u_B) - omega_B = screw_angular(u_B) + v_B = wp.spatial_top(u_B) + omega_B = wp.spatial_bottom(u_B) # Define the joint rotation offset q_x_j = Q_X_J @@ -110,7 +110,7 @@ def _set_joint_follower_body_state( # Offset the bose of the body by a fixed amount state_body_q_i[bid_F] = wp.transformation(r_F_new, q_F_new, dtype=wp.float32) - state_body_u_i[bid_F] = screw(v_F_new, omega_F_new) + state_body_u_i[bid_F] = wp.spatial_vectorf(*v_F_new, *omega_F_new) ### diff --git a/newton/_src/solvers/kamino/tests/test_kinematics_resets.py b/newton/_src/solvers/kamino/tests/test_kinematics_resets.py index de99d7fc94..5e3a17fc20 100644 --- a/newton/_src/solvers/kamino/tests/test_kinematics_resets.py +++ b/newton/_src/solvers/kamino/tests/test_kinematics_resets.py @@ -602,7 +602,7 @@ def tearDown(self): def test_01_reset_joint_states_from_body_state(self): """ - Validate that reset_joints_state_from_bodies_state() against compute_joints_data() + Validate reset_joints_state_from_bodies_state() against compute_joints_data() on a model with all joint types. """ # Initialize rng diff --git a/newton/_src/solvers/kamino/tests/test_linalg_solve_cg.py b/newton/_src/solvers/kamino/tests/test_linalg_solve_cg.py index 797d23078e..cdc5ce0021 100644 --- a/newton/_src/solvers/kamino/tests/test_linalg_solve_cg.py +++ b/newton/_src/solvers/kamino/tests/test_linalg_solve_cg.py @@ -75,7 +75,7 @@ def _test_solve(self, solver_cls, problem_params, device): maxiter=maxiter, Mi=None, callback=None, - use_cuda_graph=False, + use_graph=False, ) cur_iter, r_norm_sq, atol_sq = solver.solve(b_wp, x_wp) @@ -139,6 +139,98 @@ def test_solve_cr_cuda(self): with self.subTest(problem=problem_name, solver=solver_cls.__name__): self._test_solve(solver_cls, problem_params, device) + def _test_capture_replay_matches_eager(self, solver_cls, device): + """Regression: capture-and-replay must match an eager solve. + + Guards against the class of bug where the capture-safe path in + ``_run_capturable_loop`` silently under-iterates. Concretely, if the + break decision is frozen at record time via a host readback of stale + pre-capture memory, the recorded graph bakes in a fixed cycle count + and ``cur_iter`` under capture will be far below the eager count. + """ + device = wp.get_device(device) + problem = RandomProblemLLT( + maxdims=8, + dims=[5, 8], + seed=self.seed, + np_dtype=np.float32, + wp_dtype=wp.float32, + device=device, + ) + n_worlds = problem.num_blocks + + info = DenseSquareMultiLinearInfo() + info.finalize(dimensions=problem.maxdims, dtype=wp.float32, device=device) + info.dim = problem.dim_wp + operator = DenseLinearOperatorData(info=info, mat=problem.A_wp) + A = BatchedLinearOperator.from_dense(operator) + + world_active = wp.full(n_worlds, True, dtype=wp.bool, device=device) + maxdim = max(problem.maxdims) + atol = wp.full(n_worlds, 1.0e-4, dtype=problem.wp_dtype, device=device) + rtol = wp.full(n_worlds, 1.0e-5, dtype=problem.wp_dtype, device=device) + maxiter = wp.full(n_worlds, max(3 * maxdim, 50), dtype=int, device=device) + + def new_solver(*, use_graph): + return solver_cls( + A=A, + world_active=world_active, + atol=atol, + rtol=rtol, + maxiter=maxiter, + Mi=None, + callback=None, + use_graph=use_graph, + ) + + # Eager reference. + x_eager = wp.zeros(info.total_vec_size, dtype=wp.float32, device=device) + eager_cur, _, _ = new_solver(use_graph=False).solve(problem.b_wp, x_eager) + + # Captured replay. Warm up outside capture so kernels compile before + # recording; zero the output between the warmup and the capture body + # so both solves see the same initial state. + solver = new_solver(use_graph=True) + x_capture = wp.zeros(info.total_vec_size, dtype=wp.float32, device=device) + solver.solve(problem.b_wp, x_capture) + x_capture.zero_() + # Warp CPU graph capture currently rejects wp.copy() on non-contiguous + # arrays. Both the buggy pre-fix eager path (via strided + # ``dot_partial_sums[:, :, 0]`` in ``dot_product``) and the fixed + # capture path (via ``rz_old.assign(rz_new)`` in ``do_iteration``) + # hit that limitation on CPU, so the parity assertions below only + # activate once Warp lifts it. The CUDA counterparts exercise the + # same gate on a path where the dot buffer is contiguous and no + # skip is needed. + try: + with wp.ScopedCapture(device) as cap: + cap_cur, _, _ = solver.solve(problem.b_wp, x_capture) + except NotImplementedError as e: + if "non-contiguous" not in str(e): + raise + self.skipTest(f"Warp graph capture limitation on {device}: {e}") + assert cap.graph is not None + wp.capture_launch(cap.graph) + + np.testing.assert_array_equal(cap_cur.numpy(), eager_cur.numpy()) + np.testing.assert_array_equal(x_capture.numpy(), x_eager.numpy()) + + def test_capture_replay_cg_cpu(self): + self._test_capture_replay_matches_eager(CGSolver, "cpu") + + def test_capture_replay_cr_cpu(self): + self._test_capture_replay_matches_eager(CRSolver, "cpu") + + def test_capture_replay_cg_cuda(self): + if not wp.get_cuda_devices(): + self.skipTest("No CUDA devices found") + self._test_capture_replay_matches_eager(CGSolver, wp.get_cuda_device()) + + def test_capture_replay_cr_cuda(self): + if not wp.get_cuda_devices(): + self.skipTest("No CUDA devices found") + self._test_capture_replay_matches_eager(CRSolver, wp.get_cuda_device()) + def _test_sparse_solve(self, solver_cls, dims, block_size, device): """Test CG/CR with sparse matrices built from random SPD matrices. @@ -223,7 +315,7 @@ def _test_sparse_solve(self, solver_cls, dims, block_size, device): maxiter=None, Mi=None, callback=None, - use_cuda_graph=False, + use_graph=False, ) solver_dense.solve(b_wp, x_dense) @@ -237,7 +329,7 @@ def _test_sparse_solve(self, solver_cls, dims, block_size, device): maxiter=None, Mi=None, callback=None, - use_cuda_graph=False, + use_graph=False, ) solver_sparse.solve(b_wp, x_sparse) @@ -335,7 +427,7 @@ def test_sparse_cg_solve_simple(self): rtol=rtol, maxiter=None, Mi=None, - use_cuda_graph=False, + use_graph=False, ) solver.solve(b_wp, x_wp) @@ -477,7 +569,7 @@ def _test_solve_heterogeneous(self, solver_cls, problem_params, device): maxiter=maxiter, Mi=None, callback=None, - use_cuda_graph=False, + use_graph=False, ) solver.solve(b, x_wp) @@ -570,7 +662,7 @@ def _test_solve_heterogeneous_jacobi(self, solver_cls, problem_params, device): maxiter=maxiter, Mi=Mi, callback=None, - use_cuda_graph=False, + use_graph=False, ) solver.solve(b, x_wp) diff --git a/newton/_src/solvers/kamino/tests/test_linalg_solver_llt_blocked.py b/newton/_src/solvers/kamino/tests/test_linalg_solver_llt_blocked.py index e6426bd432..c6c839611b 100644 --- a/newton/_src/solvers/kamino/tests/test_linalg_solver_llt_blocked.py +++ b/newton/_src/solvers/kamino/tests/test_linalg_solver_llt_blocked.py @@ -50,6 +50,22 @@ def test_00_make_default_solver(self): self.assertIsNone(llt._operator) self.assertEqual(llt.dtype, wp.float32) self.assertEqual(llt.device, self.default_device) + self.assertEqual(llt._factorize_block_size, 32) + self.assertEqual(llt._solve_block_size, 32) + self.assertEqual(llt._factorize_block_dim, 128) + self.assertEqual(llt._solve_block_dim, 128) + + split_blocks = LLTBlockedSolver( + factorize_block_size=64, + solve_block_size=16, + factorize_block_dim=64, + solve_block_dim=256, + device=self.default_device, + ) + self.assertEqual(split_blocks._factorize_block_size, 64) + self.assertEqual(split_blocks._solve_block_size, 16) + self.assertEqual(split_blocks._factorize_block_dim, 64) + self.assertEqual(split_blocks._solve_block_dim, 256) def test_01_single_problem_dims_all_active(self): """ diff --git a/newton/_src/solvers/kamino/tests/test_linalg_solver_llt_blocked_rcm.py b/newton/_src/solvers/kamino/tests/test_linalg_solver_llt_blocked_rcm.py index 2f951be12f..936cf6a845 100644 --- a/newton/_src/solvers/kamino/tests/test_linalg_solver_llt_blocked_rcm.py +++ b/newton/_src/solvers/kamino/tests/test_linalg_solver_llt_blocked_rcm.py @@ -9,6 +9,7 @@ import warp as wp from newton._src.solvers.kamino._src.linalg.core import DenseLinearOperatorData, DenseSquareMultiLinearInfo +from newton._src.solvers.kamino._src.linalg.factorize import rcm_batch from newton._src.solvers.kamino._src.linalg.factorize.llt_blocked_rcm_solver import LLTBlockedRCMSolver from newton._src.solvers.kamino._src.utils import logger as msg from newton._src.solvers.kamino.tests import setup_tests, test_context @@ -50,6 +51,231 @@ def test_00_make_default_solver(self): self.assertIsNone(llt._operator) self.assertEqual(llt.dtype, wp.float32) self.assertEqual(llt.device, self.default_device) + self.assertTrue(llt._reuse_permutation) + + @staticmethod + def _run_rcm(matrix: np.ndarray, device, use_cuda_graph: bool = False) -> np.ndarray: + n = matrix.shape[0] + matrix_wp = wp.array(matrix.reshape(-1), dtype=wp.float32, device=device) + dims = wp.array([n], dtype=wp.int32, device=device) + offsets = wp.array([0], dtype=wp.int32, device=device) + permutation = wp.zeros(n, dtype=wp.int32, device=device) + scratch = rcm_batch.allocate_rcm_batch_scratch(n, 1, device) + reorder = rcm_batch.create_rcm_batch_launch( + A_flat=matrix_wp, + perm_flat=permutation, + dims=dims, + mio=offsets, + vio=offsets, + scratch=scratch, + num_blocks=1, + max_dim=n, + use_cuda_graph=use_cuda_graph, + device=device, + ) + reorder() + return permutation.numpy() + + @staticmethod + def _path_bandwidth(permutation: np.ndarray, paths: tuple[np.ndarray, ...]) -> int: + positions = np.empty(permutation.size, dtype=np.int64) + positions[permutation] = np.arange(permutation.size) + return max(int(np.max(np.abs(positions[path[:-1]] - positions[path[1:]]))) for path in paths) + + def test_complete_rcm_across_legacy_boundary(self): + """Traverse long paths completely across the former 1024-row boundary.""" + if not self.default_device.is_cuda: + self.skipTest("The legacy boundary applied only to the CUDA fast path") + + for n in (1024, 1025): + with self.subTest(n=n): + rng = np.random.default_rng(self.seed + n) + path = rng.permutation(n) + matrix = np.eye(n, dtype=np.float32) * 2.0 + matrix[path[:-1], path[1:]] = -0.25 + matrix[path[1:], path[:-1]] = -0.25 + + permutation = self._run_rcm(matrix, self.default_device, use_cuda_graph=True) + + np.testing.assert_array_equal(np.sort(permutation), np.arange(n)) + self.assertEqual(self._path_bandwidth(permutation, (path,)), 1) + + def test_complete_rcm_on_disconnected_components(self): + """Traverse every disconnected path component on each backend.""" + devices = [wp.get_device("cpu")] + if self.default_device.is_cuda: + devices.append(self.default_device) + + n = 96 + rng = np.random.default_rng(self.seed) + labels = rng.permutation(n) + paths = (labels[: n // 2], labels[n // 2 :]) + matrix = np.eye(n, dtype=np.float32) * 2.0 + for path in paths: + matrix[path[:-1], path[1:]] = -0.25 + matrix[path[1:], path[:-1]] = -0.25 + + for device in devices: + with self.subTest(device=device): + permutation = self._run_rcm(matrix, device, use_cuda_graph=device.is_cuda) + + np.testing.assert_array_equal(np.sort(permutation), np.arange(n)) + self.assertEqual(self._path_bandwidth(permutation, paths), 1) + + def test_cached_permutation_with_changed_sparsity(self): + """Verify cached RCM remains correct when numeric sparsity changes.""" + n = 96 + rng = np.random.default_rng(self.seed) + + def make_banded_spd(width): + matrix = np.zeros((n, n), dtype=np.float32) + for offset in range(1, width + 1): + values = rng.uniform(-0.2, 0.2, n - offset).astype(np.float32) + rows = np.arange(n - offset) + matrix[rows, rows + offset] = values + matrix[rows + offset, rows] = values + matrix[np.diag_indices(n)] = np.sum(np.abs(matrix), axis=1) + 1.0 + return matrix + + matrix_1 = make_banded_spd(2) + matrix_2 = make_banded_spd(9) + rhs_np = rng.standard_normal(n).astype(np.float32) + + info = DenseSquareMultiLinearInfo() + info.finalize(dimensions=[n], dtype=wp.float32, device=self.default_device) + matrix_wp = wp.array(matrix_1.reshape(-1), dtype=wp.float32, device=self.default_device) + rhs_wp = wp.array(rhs_np, dtype=wp.float32, device=self.default_device) + result_wp = wp.zeros(n, dtype=wp.float32, device=self.default_device) + operator = DenseLinearOperatorData(info=info, mat=matrix_wp) + solver = LLTBlockedRCMSolver( + operator=operator, + block_size=32, + factorize_block_dim=256, + reuse_permutation=True, + parallel_factorization=True, + device=self.default_device, + ) + + solver.compute(matrix_wp) + permutation_1 = solver.P.numpy() + matrix_wp.assign(matrix_2.reshape(-1)) + solver.compute(matrix_wp) + solver.solve(rhs_wp, result_wp) + + np.testing.assert_array_equal(solver.P.numpy(), permutation_1) + expected = np.linalg.solve(matrix_2, rhs_np) + np.testing.assert_allclose(result_wp.numpy(), expected, rtol=1.0e-3, atol=1.0e-4) + + def test_parallel_factorization_with_partial_tile(self): + """Factorize and solve a system whose final tile is partial.""" + n = 33 + matrix = np.eye(n, dtype=np.float32) * 2.0 + indices = np.arange(n - 1) + matrix[indices, indices + 1] = -0.25 + matrix[indices + 1, indices] = -0.25 + rhs = np.ones(n, dtype=np.float32) + + info = DenseSquareMultiLinearInfo() + info.finalize(dimensions=[n], dtype=wp.float32, device=self.default_device) + matrix_wp = wp.array(matrix.reshape(-1), dtype=wp.float32, device=self.default_device) + rhs_wp = wp.array(rhs, dtype=wp.float32, device=self.default_device) + result_wp = wp.zeros(n, dtype=wp.float32, device=self.default_device) + solver = LLTBlockedRCMSolver( + operator=DenseLinearOperatorData(info=info, mat=matrix_wp), + block_size=32, + reuse_permutation=True, + parallel_factorization=True, + device=self.default_device, + ) + + solver.compute(matrix_wp) + solver.solve(rhs_wp, result_wp) + + expected = np.linalg.solve(matrix, rhs) + np.testing.assert_allclose(result_wp.numpy(), expected, rtol=1.0e-4, atol=1.0e-5) + + def test_cached_permutation_on_cpu_fallback(self): + """Verify the CPU fallback reuses a cached permutation.""" + device = wp.get_device("cpu") + n = 8 + + def make_spd(edges): + matrix = np.zeros((n, n), dtype=np.float32) + for row, col in edges: + matrix[row, col] = -0.2 + matrix[col, row] = -0.2 + matrix[np.diag_indices(n)] = np.sum(np.abs(matrix), axis=1) + 1.0 + return matrix + + path_matrix = make_spd([(i, i + 1) for i in range(n - 1)]) + star_matrix = make_spd([(0, i) for i in range(1, n)]) + matrix_wp = wp.array(path_matrix.reshape(-1), dtype=wp.float32, device=device) + dims = wp.array([n], dtype=wp.int32, device=device) + offsets = wp.array([0], dtype=wp.int32, device=device) + permutation = wp.zeros(n, dtype=wp.int32, device=device) + scratch = rcm_batch.allocate_rcm_batch_scratch(n, 1, device) + reorder = rcm_batch.create_rcm_batch_launch( + A_flat=matrix_wp, + perm_flat=permutation, + dims=dims, + mio=offsets, + vio=offsets, + scratch=scratch, + num_blocks=1, + max_dim=n, + use_cuda_graph=False, + reuse_permutation=True, + device=device, + ) + + reorder() + cached = permutation.numpy() + matrix_wp.assign(star_matrix.reshape(-1)) + reorder() + np.testing.assert_array_equal(permutation.numpy(), cached) + + recomputed = wp.zeros_like(permutation) + fresh_scratch = rcm_batch.allocate_rcm_batch_scratch(n, 1, device) + recompute = rcm_batch.create_rcm_batch_launch( + A_flat=matrix_wp, + perm_flat=recomputed, + dims=dims, + mio=offsets, + vio=offsets, + scratch=fresh_scratch, + num_blocks=1, + max_dim=n, + use_cuda_graph=False, + device=device, + ) + recompute() + self.assertFalse(np.array_equal(recomputed.numpy(), cached)) + + def test_solve_with_fewer_threads_than_tile_rows(self): + """Verify the solve gathers every right-hand-side row.""" + n = 65 + rng = np.random.default_rng(self.seed) + dense = rng.standard_normal((n, n)).astype(np.float32) + matrix = dense @ dense.T + np.eye(n, dtype=np.float32) + rhs = rng.standard_normal(n).astype(np.float32) + + info = DenseSquareMultiLinearInfo() + info.finalize(dimensions=[n], dtype=wp.float32, device=self.default_device) + matrix_wp = wp.array(matrix.reshape(-1), dtype=wp.float32, device=self.default_device) + rhs_wp = wp.array(rhs, dtype=wp.float32, device=self.default_device) + result_wp = wp.zeros(n, dtype=wp.float32, device=self.default_device) + solver = LLTBlockedRCMSolver( + operator=DenseLinearOperatorData(info=info, mat=matrix_wp), + block_size=64, + solve_block_dim=32, + device=self.default_device, + ) + + solver.compute(matrix_wp) + solver.solve(rhs_wp, result_wp) + + expected = np.linalg.solve(matrix, rhs) + np.testing.assert_allclose(result_wp.numpy(), expected, rtol=1.0e-3, atol=1.0e-4) def test_01_single_problem_dims_all_active(self): """ diff --git a/newton/_src/solvers/kamino/tests/test_solver_kamino.py b/newton/_src/solvers/kamino/tests/test_solver_kamino.py index 87f3693ed5..3a77d12ae1 100644 --- a/newton/_src/solvers/kamino/tests/test_solver_kamino.py +++ b/newton/_src/solvers/kamino/tests/test_solver_kamino.py @@ -9,6 +9,7 @@ import numpy as np import warp as wp +import newton import newton._src.solvers.kamino.config as kamino_config from newton._src.solvers.kamino._src.core.control import ControlKamino from newton._src.solvers.kamino._src.core.data import DataKamino @@ -32,6 +33,7 @@ from newton._src.solvers.kamino.solver_kamino import SolverKamino from newton._src.solvers.kamino.tests import setup_tests, test_context from newton._src.solvers.kamino.tests.utils.sampling import sample_world_mask +from newton.tests.utils import basics ### # Module configs @@ -378,6 +380,59 @@ def test_01_make_explicit(self): self.assertEqual(config.padmm.warmstart_mode, "internal") +class TestCollisionCapacityInitialization(unittest.TestCase): + def setUp(self): + if not test_context.setup_done: + setup_tests(clear_cache=False) + self.default_device = wp.get_device(test_context.device) + + def _make_three_world_model(self) -> newton.Model: + source_builder = newton.ModelBuilder(up_axis=newton.Axis.Z) + SolverKamino.register_custom_attributes(source_builder) + basics.build_sphere_on_plane(builder=source_builder, z_offset=0.5) + + builder = newton.ModelBuilder(up_axis=newton.Axis.Z) + SolverKamino.register_custom_attributes(builder) + builder.replicate(source_builder, world_count=3) + return builder.finalize(device=self.default_device, skip_validation_joints=True) + + def test_capacity_allocation_for_pipeline_allocated_after_kamino(self): + """Verify capacity allocation for a pipeline allocated after Kamino.""" + model = self._make_three_world_model() + + solver = SolverKamino(model, config=SolverKamino.Config(use_collision_detector=False)) + + pipeline = newton.CollisionPipeline(model) + + # Kamino's capacity should be equal to the pipeline's capacity. + self.assertEqual(solver._contacts_kamino.model_max_contacts_host, pipeline.rigid_contact_max) + + def test_capacity_allocation_for_pipeline_allocated_before_kamino(self): + """Verify capacity allocation for a pipeline allocated before Kamino.""" + model = self._make_three_world_model() + + pipeline = newton.CollisionPipeline(model) + solver = SolverKamino(model, config=SolverKamino.Config(use_collision_detector=False)) + + # Kamino's capacity should be at least as large as the pipeline's capacity. Do to rounding up to the nearest multiple of the world count. + self.assertGreaterEqual(solver._contacts_kamino.model_max_contacts_host, pipeline.rigid_contact_max) + + def test_external_collisions_preserve_explicit_rigid_contact_max(self): + """Verify external collisions preserve an explicit contact capacity.""" + model = self._make_three_world_model() + model.rigid_contact_max = 1000 + + solver = SolverKamino(model, config=SolverKamino.Config(use_collision_detector=False)) + + self.assertEqual(model.rigid_contact_max, 1000) + self.assertEqual(solver._contacts_kamino.world_max_contacts_host, [334, 334, 334]) + self.assertEqual(solver._contacts_kamino.model_max_contacts_host, 1002) + + contacts = newton.CollisionPipeline(model).contacts() + with self.assertNoLogs(level="WARNING"): + solver.update_contacts(contacts, model.state()) + + class TestSolverKaminoImpl(unittest.TestCase): def setUp(self): if not test_context.setup_done: diff --git a/newton/_src/solvers/kamino/tests/test_solver_kamino_notify.py b/newton/_src/solvers/kamino/tests/test_solver_kamino_notify.py index dff362a68e..a543049fd0 100644 --- a/newton/_src/solvers/kamino/tests/test_solver_kamino_notify.py +++ b/newton/_src/solvers/kamino/tests/test_solver_kamino_notify.py @@ -12,6 +12,7 @@ import warp as wp import newton +from newton._src.solvers.kamino._src.core.materials import DEFAULT_FRICTION, DEFAULT_RESTITUTION from newton._src.solvers.kamino.solver_kamino import SolverKamino from newton._src.solvers.kamino.tests import setup_tests, test_context @@ -22,10 +23,14 @@ def _build_revolute( limited: bool = False, actuator_mode: newton.JointTargetMode = newton.JointTargetMode.NONE, body_com: wp.vec3f | None = None, + shape_materials: tuple[tuple[float, float], ...] | None = None, + has_shape_collision: bool = True, + fk_actuation_flag: int | None = None, ) -> newton.Model: """Build a tiny world-to-body revolute model for notify tests.""" builder = newton.ModelBuilder() - SolverKamino.register_custom_attributes(builder) + fk_actuation_flags = None if fk_actuation_flag is None else {0: fk_actuation_flag} + SolverKamino.register_custom_attributes(builder, fk_actuation_flags=fk_actuation_flags) builder.begin_world() bid = builder.add_link( @@ -36,7 +41,33 @@ def _build_revolute( com=body_com, lock_inertia=True, ) - builder.add_shape_box(label="box", body=bid, hx=0.1, hy=0.1, hz=0.1) + if shape_materials is None: + builder.add_shape_box( + label="box", + body=bid, + hx=0.1, + hy=0.1, + hz=0.1, + cfg=newton.ModelBuilder.ShapeConfig(has_shape_collision=has_shape_collision), + ) + else: + for shape, (mu, restitution) in enumerate(shape_materials): + builder.add_shape_box( + label=f"box_{shape}", + body=bid, + xform=wp.transformf( + wp.vec3f(0.3 * shape, 0.0, 0.0), + wp.quat_identity(dtype=wp.float32), + ), + hx=0.1, + hy=0.1, + hz=0.1, + cfg=newton.ModelBuilder.ShapeConfig( + mu=mu, + restitution=restitution, + has_shape_collision=has_shape_collision, + ), + ) jid = builder.add_joint_revolute( label="world_to_link", @@ -58,6 +89,42 @@ def _build_revolute( return builder.finalize() +def _build_free_body() -> newton.Model: + """Build one free body so FK creates a synthetic base joint.""" + builder = newton.ModelBuilder() + SolverKamino.register_custom_attributes(builder) + builder.begin_world() + bid = builder.add_link( + label="base", + mass=1.0, + inertia=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], + xform=wp.transformf(wp.vec3f(0.0, 0.0, 1.0), wp.quat_identity(dtype=wp.float32)), + lock_inertia=True, + ) + builder.add_shape_box(body=bid, hx=0.1, hy=0.1, hz=0.1) + builder.end_world() + return builder.finalize() + + +def _build_free_root(*, fk_actuation_flag: int = -1) -> newton.Model: + """Build one body attached to the world by an explicit free root joint.""" + builder = newton.ModelBuilder() + SolverKamino.register_custom_attributes(builder, fk_actuation_flags={0: fk_actuation_flag}) + builder.begin_world() + bid = builder.add_link( + label="base", + mass=1.0, + inertia=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], + xform=wp.transformf(wp.vec3f(0.0, 0.0, 1.0), wp.quat_identity(dtype=wp.float32)), + lock_inertia=True, + ) + builder.add_shape_box(body=bid, hx=0.1, hy=0.1, hz=0.1) + jid = builder.add_joint_free(parent=-1, child=bid) + builder.add_articulation([jid]) + builder.end_world() + return builder.finalize() + + def _snapshot_model_arrays(model: newton.Model) -> dict[str, np.ndarray]: """Copy every allocated top-level Warp array on a model.""" return {name: value.numpy().copy() for name, value in vars(model).items() if isinstance(value, wp.array)} @@ -85,6 +152,7 @@ def test_noop_flags_are_silent_and_do_not_mutate_newton_arrays(self): solver = SolverKamino(model) snapshot = _snapshot_model_arrays(model) noop_flags = ( + newton.ModelFlags.MODEL_PROPERTIES, newton.ModelFlags.BODY_PROPERTIES, newton.ModelFlags.BODY_INERTIAL_PROPERTIES, newton.ModelFlags.SHAPE_PROPERTIES, @@ -130,9 +198,11 @@ def test_aliased_properties_reference_newton(self): bodies = solver._model_kamino.bodies joints = solver._model_kamino.joints geoms = solver._model_kamino.geoms + gravity = solver._model_kamino.gravity # (Newton model attribute, Kamino container, Kamino attribute) for each direct alias. aliased_properties = [ + ("gravity", gravity, "vector"), ("body_mass", bodies, "m_i"), ("body_inv_mass", bodies, "inv_m_i"), ("body_com", bodies, "i_r_com_i"), @@ -168,21 +238,6 @@ def test_aliased_properties_reference_newton(self): newton_array.assign(perturbed) np.testing.assert_array_equal(kamino_array.numpy(), perturbed) - def test_gravity_update(self): - """Model-property notifications refresh Kamino's gravity representation.""" - model = _build_revolute(limited=True) - solver = SolverKamino(model) - gravity = np.tile(np.array([1.0, -2.0, 3.0], dtype=np.float32), (model.world_count, 1)) - acceleration = np.linalg.norm(gravity, axis=1) - - model.gravity.assign(gravity) - solver.notify_model_changed(newton.ModelFlags.MODEL_PROPERTIES) - - expected_g_dir_acc = np.column_stack((gravity / acceleration[:, None], acceleration)) - expected_vector = np.column_stack((gravity, np.ones(model.world_count, dtype=np.float32))) - np.testing.assert_allclose(solver._model_kamino.gravity.g_dir_acc.numpy(), expected_g_dir_acc, atol=1e-6) - np.testing.assert_allclose(solver._model_kamino.gravity.vector.numpy(), expected_vector, atol=1e-6) - def test_joint_transform_update(self): """Joint-property notifications recompute Kamino's parent and child frames.""" model = _build_revolute(limited=True) @@ -295,6 +350,142 @@ def test_body_com_refreshes_derived_quantities(self): solver.reset(state) np.testing.assert_allclose(state.body_q.numpy(), model.body_q.numpy(), atol=1e-6) + def test_material_value_update_propagates(self): + """Two shapes sharing one material can update it together and keep sharing it.""" + model = _build_revolute(shape_materials=((0.2, 0.1), (0.2, 0.1))) + solver = SolverKamino(model, SolverKamino.Config(use_collision_detector=True)) + materials = solver._model_kamino.materials + material_pairs = solver._model_kamino.material_pairs + arrays = ( + materials.restitution, + materials.static_friction, + materials.dynamic_friction, + material_pairs.restitution, + material_pairs.static_friction, + material_pairs.dynamic_friction, + ) + pointers = tuple(array.ptr for array in arrays) + pair_values = tuple(array.numpy().copy() for array in arrays[3:]) + + model.shape_material_mu.assign(np.array([0.4, 0.4], dtype=np.float32)) + model.shape_material_restitution.assign(np.array([0.3, 0.3], dtype=np.float32)) + solver.notify_model_changed(newton.ModelFlags.SHAPE_PROPERTIES) + + self.assertEqual(tuple(array.ptr for array in arrays), pointers) + np.testing.assert_allclose(materials.static_friction.numpy(), [DEFAULT_FRICTION, 0.4]) + np.testing.assert_allclose(materials.dynamic_friction.numpy(), [DEFAULT_FRICTION, 0.4]) + np.testing.assert_allclose(materials.restitution.numpy(), [DEFAULT_RESTITUTION, 0.3]) + for actual, expected in zip(arrays[3:], pair_values, strict=True): + np.testing.assert_array_equal(actual.numpy(), expected) + + def test_default_material_update_propagates_to_default_pair(self): + """Updating material zero keeps its explicit self-pair synchronized.""" + model = _build_revolute(shape_materials=((DEFAULT_FRICTION, DEFAULT_RESTITUTION),)) + solver = SolverKamino(model, SolverKamino.Config(use_collision_detector=True)) + materials = solver._model_kamino.materials + material_pairs = solver._model_kamino.material_pairs + np.testing.assert_array_equal(solver._model_kamino.geoms.material.numpy(), [0]) + + model.shape_material_mu.assign([0.4]) + model.shape_material_restitution.assign([0.3]) + solver.notify_model_changed(newton.ModelFlags.SHAPE_PROPERTIES) + + np.testing.assert_allclose(materials.static_friction.numpy(), [0.4]) + np.testing.assert_allclose(materials.dynamic_friction.numpy(), [0.4]) + np.testing.assert_allclose(materials.restitution.numpy(), [0.3]) + np.testing.assert_allclose(material_pairs.static_friction.numpy(), [0.4]) + np.testing.assert_allclose(material_pairs.dynamic_friction.numpy(), [0.4]) + np.testing.assert_allclose(material_pairs.restitution.numpy(), [0.3]) + + def test_material_ids_can_converge_to_same_values(self): + """Distinct material IDs remain valid when their coefficients become equal.""" + model = _build_revolute(shape_materials=((0.2, 0.1), (0.6, 0.5))) + solver = SolverKamino(model, SolverKamino.Config(use_collision_detector=True)) + materials = solver._model_kamino.materials + geoms = solver._model_kamino.geoms + material_mapping = geoms.material.numpy().copy() + + model.shape_material_mu.assign(np.array([0.4, 0.4], dtype=np.float32)) + model.shape_material_restitution.assign(np.array([0.3, 0.3], dtype=np.float32)) + solver.notify_model_changed(newton.ModelFlags.SHAPE_PROPERTIES) + + np.testing.assert_array_equal(geoms.material.numpy(), material_mapping) + np.testing.assert_allclose(materials.static_friction.numpy(), [DEFAULT_FRICTION, 0.4, 0.4]) + np.testing.assert_allclose(materials.dynamic_friction.numpy(), [DEFAULT_FRICTION, 0.4, 0.4]) + np.testing.assert_allclose(materials.restitution.numpy(), [DEFAULT_RESTITUTION, 0.3, 0.3]) + + def test_shape_without_material_is_ignored(self): + """Shapes without a Kamino material mapping do not modify material tables.""" + model = _build_revolute(shape_materials=((0.2, 0.1),), has_shape_collision=False) + solver = SolverKamino(model) + materials = solver._model_kamino.materials + before = ( + materials.restitution.numpy().copy(), + materials.static_friction.numpy().copy(), + materials.dynamic_friction.numpy().copy(), + ) + # Non-collidable shapes use -1 to indicate that they need no contact material. + np.testing.assert_array_equal(solver._model_kamino.geoms.material.numpy(), [-1]) + model.shape_material_mu.assign([0.7]) + model.shape_material_restitution.assign([0.8]) + + solver.notify_model_changed(newton.ModelFlags.SHAPE_PROPERTIES) + + for actual, expected in zip( + (materials.restitution, materials.static_friction, materials.dynamic_friction), + before, + strict=True, + ): + np.testing.assert_array_equal(actual.numpy(), expected) + + def test_material_structural_change_raises(self): + """Two shapes sharing one material cannot update it to different values.""" + model = _build_revolute(shape_materials=((0.2, 0.1), (0.2, 0.1))) + solver = SolverKamino(model, SolverKamino.Config(use_collision_detector=True)) + materials = solver._model_kamino.materials + before = ( + materials.restitution.numpy().copy(), + materials.static_friction.numpy().copy(), + materials.dynamic_friction.numpy().copy(), + ) + model.shape_material_mu.assign(np.array([0.2, 0.4], dtype=np.float32)) + + with self.assertRaisesRegex(RuntimeError, "recreate"): + solver.notify_model_changed(newton.ModelFlags.SHAPE_PROPERTIES) + + for actual, expected in zip( + (materials.restitution, materials.static_friction, materials.dynamic_friction), + before, + strict=True, + ): + np.testing.assert_array_equal(actual.numpy(), expected) + + def test_external_collisions_allow_material_structural_change(self): + """Allow per-shape material changes when using external Newton collisions.""" + model = _build_revolute(shape_materials=((0.2, 0.1), (0.2, 0.1))) + solver = SolverKamino(model, SolverKamino.Config(use_collision_detector=False)) + materials = solver._model_kamino.materials + before = ( + materials.restitution.numpy().copy(), + materials.static_friction.numpy().copy(), + materials.dynamic_friction.numpy().copy(), + ) + updated_friction = np.array([0.2, 0.4], dtype=np.float32) + updated_restitution = np.array([0.1, 0.3], dtype=np.float32) + model.shape_material_mu.assign(updated_friction) + model.shape_material_restitution.assign(updated_restitution) + + solver.notify_model_changed(newton.ModelFlags.SHAPE_PROPERTIES) + + np.testing.assert_array_equal(model.shape_material_mu.numpy(), updated_friction) + np.testing.assert_array_equal(model.shape_material_restitution.numpy(), updated_restitution) + for actual, expected in zip( + (materials.restitution, materials.static_friction, materials.dynamic_friction), + before, + strict=True, + ): + np.testing.assert_array_equal(actual.numpy(), expected) + def test_dynamic_constraint_toggle_raises(self): """Adding or removing a joint's dynamic constraints requires solver recreation.""" for built_dynamic in (False, True): @@ -387,6 +578,188 @@ def test_active_actuation_mode_change_is_allowed(self): expected = solver._kamino.JointActuationType.from_newton(changed_mode) self.assertEqual(solver._model_kamino.joints.act_type.numpy()[0], expected) + def test_fk_joint_frame_changes_propagate(self): + """Joint and CoM notifications propagate to FK-owned frames.""" + model = _build_revolute(actuator_mode=newton.JointTargetMode.POSITION) + solver = SolverKamino( + model, + SolverKamino.Config(use_fk_solver=True, use_collision_detector=False), + ) + fk = solver._solver_kamino.solver_fk + + model.joint_X_p.assign( + [wp.transformf(wp.vec3f(0.2, -0.1, 0.3), wp.quat_from_axis_angle(wp.vec3f(0.0, 0.0, 1.0), 0.4))] + ) + model.joint_X_c.assign( + [wp.transformf(wp.vec3f(-0.4, 0.5, 0.6), wp.quat_from_axis_angle(wp.vec3f(1.0, 0.0, 0.0), -0.35))] + ) + solver.notify_model_changed(newton.ModelFlags.JOINT_PROPERTIES) + + model.body_com.assign([wp.vec3f(0.1, -0.2, 0.15)]) + solver.notify_model_changed(newton.ModelFlags.BODY_INERTIAL_PROPERTIES) + + fk_joint = int(np.flatnonzero(fk.joints_source_id.numpy() == 0)[0]) + joints = solver._model_kamino.joints + for fk_values, model_values in ( + (fk.joints_B_r_Bj, joints.B_r_Bj), + (fk.joints_F_r_Fj, joints.F_r_Fj), + (fk.joints_X_Bj, joints.X_Bj), + (fk.joints_X_Fj, joints.X_Fj), + ): + np.testing.assert_allclose(fk_values.numpy()[fk_joint], model_values.numpy()[0], atol=1e-6) + + def test_fk_base_pose_changes_propagate(self): + """Body-pose notifications propagate to the default synthetic FK base pose.""" + model = _build_free_body() + solver = SolverKamino( + model, + SolverKamino.Config(use_fk_solver=True, use_collision_detector=False), + ) + fk = solver._solver_kamino.solver_fk + new_pose = wp.transformf( + wp.vec3f(0.3, -0.4, 1.5), + wp.quat_from_axis_angle(wp.vec3f(0.0, 1.0, 0.0), 0.25), + ) + model.body_q.assign([new_pose]) + + solver.notify_model_changed(newton.ModelFlags.BODY_PROPERTIES) + + np.testing.assert_allclose( + fk.base_q_default.numpy()[0], + solver._model_kamino.bodies.q_i_0.numpy()[0], + atol=1e-6, + ) + + def test_fk_explicit_base_pose_changes_propagate(self): + """Joint-property notifications refresh an explicit FK base pose.""" + model = _build_free_root() + solver = SolverKamino( + model, + SolverKamino.Config(use_fk_solver=True, use_collision_detector=False), + ) + fk = solver._solver_kamino.solver_fk + new_pose = wp.transformf( + wp.vec3f(0.3, -0.4, 1.5), + wp.quat_from_axis_angle(wp.vec3f(0.0, 1.0, 0.0), 0.25), + ) + model.joint_q.assign(np.asarray(new_pose)) + + solver.notify_model_changed(newton.ModelFlags.JOINT_PROPERTIES) + + np.testing.assert_allclose( + fk.base_q_default.numpy()[0], + np.asarray(new_pose), + atol=1e-6, + ) + + def test_fk_actuation_partition_change_raises(self): + """Runtime FK override edits cannot change the FK buffer layout.""" + for flag in (newton.ModelFlags.ACTUATOR_PROPERTIES, newton.ModelFlags.JOINT_DOF_PROPERTIES): + with self.subTest(flag=flag.name): + model = _build_revolute( + actuator_mode=newton.JointTargetMode.POSITION, + fk_actuation_flag=1, + ) + solver = SolverKamino( + model, + SolverKamino.Config(use_fk_solver=True, use_collision_detector=False), + ) + model.fk_actuation_flag.assign([0]) + + with self.assertRaisesRegex(RuntimeError, "actuated vs passive status.*recreate"): + solver.notify_model_changed(flag) + + def test_fk_base_joint_override_change_is_allowed(self): + """FK overrides do not affect explicit base joints replaced by free joints.""" + for flag in (newton.ModelFlags.ACTUATOR_PROPERTIES, newton.ModelFlags.JOINT_DOF_PROPERTIES): + with self.subTest(flag=flag.name): + model = _build_free_root(fk_actuation_flag=0) + solver = SolverKamino( + model, + SolverKamino.Config(use_fk_solver=True, use_collision_detector=False), + ) + fk = solver._solver_kamino.solver_fk + model.fk_actuation_flag.assign([1]) + + solver.notify_model_changed(flag) + + self.assertEqual(fk.joints_act_type.numpy()[0], solver._kamino.JointActuationType.FORCE) + + def test_equivalent_fk_actuation_override_change_is_allowed(self): + """Raw FK override changes are allowed when effective actuation is unchanged.""" + model = _build_revolute( + actuator_mode=newton.JointTargetMode.POSITION, + fk_actuation_flag=1, + ) + solver = SolverKamino( + model, + SolverKamino.Config(use_fk_solver=True, use_collision_detector=False), + ) + fk = solver._solver_kamino.solver_fk + model.fk_actuation_flag.assign([-1]) + + solver.notify_model_changed(newton.ModelFlags.ACTUATOR_PROPERTIES) + + fk_joint = int(np.flatnonzero(fk.joints_source_id.numpy() == 0)[0]) + self.assertNotEqual( + fk.joints_act_type.numpy()[fk_joint], + solver._kamino.JointActuationType.PASSIVE, + ) + + def test_invalid_fk_actuation_override_raises(self): + """Runtime FK overrides accept only the documented -1, 0, and 1 values.""" + models = ( + _build_revolute( + actuator_mode=newton.JointTargetMode.POSITION, + fk_actuation_flag=1, + ), + _build_free_root(fk_actuation_flag=0), + ) + for model in models: + with self.subTest(joint_type=newton.JointType(model.joint_type.numpy()[0]).name): + solver = SolverKamino( + model, + SolverKamino.Config(use_fk_solver=True, use_collision_detector=False), + ) + model.fk_actuation_flag.assign([2]) + + with self.assertRaisesRegex(ValueError, "Invalid FK actuation flag"): + solver.notify_model_changed(newton.ModelFlags.ACTUATOR_PROPERTIES) + + def test_fk_reset_matches_fresh_solver_after_joint_update(self): + """An FK reset after notify matches a solver built from the updated model.""" + model = _build_revolute(actuator_mode=newton.JointTargetMode.POSITION) + config = SolverKamino.Config(use_fk_solver=True, use_collision_detector=False) + solver = SolverKamino(model, config) + model.joint_X_c.assign( + [wp.transformf(wp.vec3f(0.2, 0.1, -0.15), wp.quat_from_axis_angle(wp.vec3f(1.0, 0.0, 0.0), 0.2))] + ) + solver.notify_model_changed(newton.ModelFlags.JOINT_PROPERTIES) + reference = SolverKamino(model, SolverKamino.Config(use_fk_solver=True, use_collision_detector=False)) + actuator_q = wp.array([0.35], dtype=wp.float32, device=model.device) + reset_config = SolverKamino.ResetConfig( + body_poses=SolverKamino.ResetConfig.FromActuatorQ(actuator_q), + ) + state = model.state() + reference_state = model.state() + + solver.reset(state, config=reset_config) + reference.reset(reference_state, config=reset_config) + + np.testing.assert_allclose(state.body_q.numpy(), reference_state.body_q.numpy(), atol=1e-5) + + def test_invalid_actuation_mode_raises_before_update(self): + """Invalid target modes do not mutate Kamino's actuation table.""" + model = _build_revolute(actuator_mode=newton.JointTargetMode.POSITION) + solver = SolverKamino(model) + before = solver._model_kamino.joints.act_type.numpy().copy() + model.joint_target_mode.assign([99]) + + with self.assertRaisesRegex(ValueError, "Unsupported joint target mode"): + solver.notify_model_changed(newton.ModelFlags.ACTUATOR_PROPERTIES) + + np.testing.assert_array_equal(solver._model_kamino.joints.act_type.numpy(), before) + if __name__ == "__main__": unittest.main() diff --git a/newton/_src/solvers/kamino/tests/test_solvers_dvi.py b/newton/_src/solvers/kamino/tests/test_solvers_dvi.py new file mode 100644 index 0000000000..939fae8cc5 --- /dev/null +++ b/newton/_src/solvers/kamino/tests/test_solvers_dvi.py @@ -0,0 +1,1448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the Kamino DVI solver.""" + +from __future__ import annotations + +import unittest +from types import SimpleNamespace + +import numpy as np +import warp as wp + +import newton +import newton._src.solvers.kamino.config as kamino_config +from newton._src.solvers.kamino._src.dynamics.dual import DualProblem +from newton._src.solvers.kamino._src.integrators.euler import integrate_euler_semi_implicit +from newton._src.solvers.kamino._src.kinematics.constraints import unpack_constraint_solutions, update_constraints_info +from newton._src.solvers.kamino._src.kinematics.jacobians import DenseSystemJacobians +from newton._src.solvers.kamino._src.linalg import LLTBlockedRCMSolver, LLTBlockedSolver +from newton._src.solvers.kamino._src.models.builders import basics, testing +from newton._src.solvers.kamino._src.models.builders import utils as builder_utils +from newton._src.solvers.kamino._src.solvers.common import WarmStartMode +from newton._src.solvers.kamino._src.solvers.dvi import DVISolver +from newton._src.solvers.kamino._src.solvers.dvi.kernels import _color_dvi_contacts, _initialize_dvi_status +from newton._src.solvers.kamino._src.solvers.dvi.sparse import ( + _SPARSE_DELASSUS_ROWS_JOINTS, + _SPARSE_DELASSUS_ROWS_UNILATERAL, + _sparse_delassus_matvec_rows, +) +from newton._src.solvers.kamino._src.solvers.metrics import SolutionMetrics +from newton._src.solvers.kamino.solver_kamino import SolverKamino +from newton._src.solvers.kamino.tests import setup_tests, test_context +from newton._src.solvers.kamino.tests.test_solvers_padmm import TestSetup +from newton._src.solvers.kamino.tests.utils.extract import extract_delassus, extract_problem_vector +from newton._src.solvers.kamino.tests.utils.make import make_containers, make_test_problem_fourbar, update_containers +from newton.tests.utils import basics as public_basics + + +def _reduce_solver_status(status: np.ndarray) -> dict[str, object]: + """Reduce per-world status while requiring every world to converge.""" + return { + name: bool(np.all(status[name])) if name == "converged" else np.max(status[name]).item() + for name in status.dtype.names + } + + +def _check_solution_matches_dual_problem(testcase: unittest.TestCase, problem: DualProblem, solver: DVISolver): + """Check that final physical solution vectors match ``D lambda + v_f``.""" + D_np = extract_delassus(problem.delassus, only_active_dims=True) + v_f_np = extract_problem_vector(problem.delassus, problem.data.v_f.numpy(), only_active_dims=True) + P_np = extract_problem_vector(problem.delassus, problem.data.P.numpy(), only_active_dims=True) + lambdas_np = extract_problem_vector(problem.delassus, solver.data.solution.lambdas.numpy(), only_active_dims=True) + v_plus_np = extract_problem_vector(problem.delassus, solver.data.solution.v_plus.numpy(), only_active_dims=True) + + status = solver.data.status.numpy() + for wid in range(problem.data.num_worlds): + P_inv = np.diag(np.reciprocal(P_np[wid])) + D_true = P_inv @ D_np[wid] @ P_inv + v_f_true = P_inv @ v_f_np[wid] + v_plus_true = D_true @ lambdas_np[wid] + v_f_true + np.testing.assert_allclose(v_plus_np[wid], v_plus_true, rtol=1e-4, atol=1e-4) + + testcase.assertEqual(int(status[wid]["converged"]), 1) + testcase.assertLessEqual(int(status[wid]["iterations"]), _status_iteration_budget(solver, wid)) + testcase.assertLessEqual(float(status[wid]["r_p"]), solver.config[wid].tolerance) + testcase.assertLessEqual(float(status[wid]["r_d"]), solver.config[wid].tolerance) + testcase.assertLessEqual(float(status[wid]["r_c"]), solver.config[wid].tolerance) + testcase.assertLessEqual(float(status[wid]["r_b"]), solver.config[wid].tolerance) + + +def _make_dense_dual_problem(model, data, limits, contacts, jacobians) -> DualProblem: + problem = DualProblem( + model=model, + data=data, + limits=limits, + contacts=contacts, + jacobians=jacobians, + solver=LLTBlockedSolver, + sparse=False, + ) + problem.build(model=model, data=data, limits=limits, contacts=contacts, jacobians=jacobians) + return problem + + +def _make_sparse_dual_problem(model, data, limits, contacts, jacobians) -> DualProblem: + problem = DualProblem( + model=model, + data=data, + limits=limits, + contacts=contacts, + jacobians=jacobians, + sparse=True, + ) + problem.build(model=model, data=data, limits=limits, contacts=contacts, jacobians=jacobians) + return problem + + +def _solve_dvi( + model, + problem, + warmstart: WarmStartMode = WarmStartMode.NONE, + config: kamino_config.DVISolverConfig | None = None, + setup: TestSetup | None = None, +) -> DVISolver: + solver = DVISolver( + model=model, + data=setup.data if setup is not None else None, + limits=setup.limits if setup is not None else None, + contacts=setup.contacts if setup is not None else None, + jacobians=setup.jacobians if setup is not None else None, + config=config or kamino_config.DVISolverConfig(max_iterations=300, tolerance=1e-4, regularization=1e-5), + warmstart=warmstart, + ) + solver.reset() + solver.coldstart() + solver.solve(problem) + return solver + + +def _status_iteration_budget(solver: DVISolver, wid: int) -> int: + config = solver.config[wid] + if solver._bilateral_solver is not None and solver.data.bilateral_operator is not None: + return max(config.max_iterations, config.block_iterations * config.contact_iterations) + return config.max_iterations + + +def _assert_solver_status_converged(testcase: unittest.TestCase, solver: DVISolver): + status = solver.data.status.numpy() + for wid in range(solver.size.num_worlds): + testcase.assertEqual(int(status[wid]["converged"]), 1) + testcase.assertLessEqual(int(status[wid]["iterations"]), _status_iteration_budget(solver, wid)) + testcase.assertLessEqual(float(status[wid]["r_p"]), solver.config[wid].tolerance) + testcase.assertLessEqual(float(status[wid]["r_d"]), solver.config[wid].tolerance) + testcase.assertLessEqual(float(status[wid]["r_c"]), solver.config[wid].tolerance) + testcase.assertLessEqual(float(status[wid]["r_b"]), solver.config[wid].tolerance) + + +def _assert_solution_finite(testcase: unittest.TestCase, solver: DVISolver): + testcase.assertTrue(np.all(np.isfinite(solver.data.solution.lambdas.numpy()))) + testcase.assertTrue(np.all(np.isfinite(solver.data.solution.v_plus.numpy()))) + + +def _evaluate_solution_metrics(test: TestSetup, solver: DVISolver) -> dict[str, float]: + integrate_euler_semi_implicit(model=test.model, data=test.data) + metrics = SolutionMetrics(model=test.model) + metrics.evaluate( + sigma=solver.data.state.sigma, + lambdas=solver.data.solution.lambdas, + v_plus=solver.data.solution.v_plus, + model=test.model, + data=test.data, + state_p=test.state_p, + problem=test.problem, + jacobians=test.jacobians, + limits=test.limits, + contacts=test.contacts, + ) + return { + name: float(np.max(getattr(metrics.data, name).numpy())) + for name in ( + "r_eom", + "r_kinematics", + "r_cts_joints", + "r_cts_limits", + "r_cts_contacts", + "r_v_plus", + "r_ncp_primal", + "r_ncp_dual", + "r_ncp_compl", + "r_vi_natmap", + ) + } + + +class TestDVISolver(unittest.TestCase): + def setUp(self): + if not test_context.setup_done: + setup_tests(clear_cache=False) + self.device = wp.get_device(test_context.device) + + def test_00_config_selection(self): + default_config = SolverKamino.Config(dynamics_solver="dvi") + self.assertFalse(default_config.sparse_dynamics) + self.assertTrue(default_config.sparse_jacobian) + self.assertEqual(default_config.integrator, "euler") + self.assertEqual(default_config.dynamics.linear_solver_type, "LLTBRCM") + self.assertEqual(default_config.dynamics.linear_solver_kwargs, {}) + self.assertEqual(default_config.dvi.omega, 1.0) + self.assertEqual(default_config.dvi.block_iterations, 32) + self.assertEqual(default_config.dvi.contact_iterations, 4) + self.assertEqual(default_config.dvi.bilateral_solve_period, 1) + self.assertEqual(default_config.dvi.bilateral_solver_type, "LLTB") + self.assertEqual(default_config.dvi.bilateral_solver_kwargs, {}) + self.assertEqual(default_config.dvi.contact_jacobi_omega, 0.3) + self.assertEqual(default_config.dvi.contact_jacobi_relaxation, 0.9) + + dense_config = SolverKamino.Config( + dynamics_solver="dvi", + sparse_dynamics=False, + sparse_jacobian=False, + ) + self.assertFalse(dense_config.sparse_dynamics) + self.assertFalse(dense_config.sparse_jacobian) + self.assertEqual(dense_config.integrator, "euler") + self.assertEqual(dense_config.dynamics.linear_solver_type, "LLTBRCM") + self.assertEqual(dense_config.dvi.block_iterations, 32) + + padmm_config = SolverKamino.Config() + self.assertFalse(padmm_config.sparse_dynamics) + self.assertFalse(padmm_config.sparse_jacobian) + self.assertEqual(padmm_config.integrator, "euler") + + config = SolverKamino.Config( + dynamics_solver="dvi", + dvi=kamino_config.DVISolverConfig(max_iterations=32, tolerance=1e-4), + ) + self.assertEqual(config.dynamics_solver, "dvi") + self.assertEqual(config.dvi.max_iterations, 32) + self.assertEqual(config.dvi.block_iterations, 32) + self.assertEqual(config.dvi.contact_iterations, 4) + self.assertEqual(config.dvi.bilateral_solve_period, 1) + self.assertEqual(config.dvi.contact_jacobi_omega, 0.3) + self.assertEqual(config.dvi.contact_jacobi_relaxation, 0.9) + self.assertFalse(config.dvi.contact_block_preconditioner) + self.assertEqual(config.dvi.contact_warmstart_method, "key_and_position_with_net_force_backup") + self.assertFalse(config.dynamics.preconditioning) + + sparse_config = SolverKamino.Config(dynamics_solver="dvi", sparse_dynamics=True, sparse_jacobian=True) + self.assertTrue(sparse_config.sparse_dynamics) + self.assertTrue(sparse_config.sparse_jacobian) + self.assertEqual(sparse_config.dynamics.linear_solver_type, "CR") + self.assertEqual(sparse_config.dynamics.linear_solver_kwargs, {"maxiter": 9}) + with self.assertRaises(ValueError): + SolverKamino.Config( + dynamics_solver="dvi", + dynamics=kamino_config.ConstrainedDynamicsConfig(preconditioning=True), + ) + invalid_dvi_configs = ( + {"max_iterations": 0}, + {"tolerance": -1.0}, + {"regularization": 0.0}, + {"omega": 0.0}, + {"omega": 2.1}, + {"block_iterations": 0}, + {"contact_iterations": 0}, + {"bilateral_solve_period": 0}, + {"bilateral_solver_type": "invalid"}, + {"contact_jacobi_omega": 0.0}, + {"contact_jacobi_omega": 2.1}, + {"contact_jacobi_relaxation": 0.0}, + {"contact_jacobi_relaxation": 1.1}, + {"warmstart_mode": "invalid"}, + ) + for kwargs in invalid_dvi_configs: + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + kamino_config.DVISolverConfig(**kwargs) + for method in ( + "key_and_position", + "geom_pair_net_force", + "key_and_position_with_net_force_backup", + ): + self.assertEqual( + kamino_config.DVISolverConfig(contact_warmstart_method=method).contact_warmstart_method, method + ) + for method in ("reaction", "geom_pair_net_wrench"): + with self.assertRaises(ValueError): + kamino_config.DVISolverConfig(contact_warmstart_method=method) + + model_with_attrs = SimpleNamespace( + kamino=SimpleNamespace(max_solver_iterations=wp.array([37], dtype=wp.int32, device=self.device)) + ) + self.assertEqual(kamino_config.DVISolverConfig.from_model(model_with_attrs).max_iterations, 37) + + def test_00b_bilateral_solver_selection(self): + """Verify DVI constructs and validates the configured bilateral solver.""" + + def make_model(dimensions): + return SimpleNamespace( + size=SimpleNamespace(sum_of_num_joint_cts=sum(dimensions)), + info=SimpleNamespace( + num_joint_cts=wp.array(dimensions, dtype=wp.int32, device=self.device), + joint_cts_offset=wp.array(np.cumsum([0, *dimensions[:-1]]), dtype=wp.int32, device=self.device), + ), + ) + + config = kamino_config.DVISolverConfig( + bilateral_solver_type="LLTBRCM", + bilateral_solver_kwargs={"block_size": 16, "reuse_permutation": True}, + ) + solver = DVISolver() + solver._config = [config] + solver._data = SimpleNamespace(bilateral_operator=None) + solver._device = self.device + solver._allocate_bilateral_solver(make_model([3])) + + self.assertIsInstance(solver._bilateral_solver, LLTBlockedRCMSolver) + self.assertEqual(solver._bilateral_solver._block_size, 16) + self.assertTrue(solver._bilateral_solver._reuse_permutation) + + solver._config = [ + kamino_config.DVISolverConfig(bilateral_solver_type="LLTB"), + kamino_config.DVISolverConfig(bilateral_solver_type="LLTBRCM"), + ] + with self.assertRaisesRegex(ValueError, "All worlds must use the same"): + solver._allocate_bilateral_solver(make_model([3, 3])) + + def test_00a_multiworld_status_reduction_requires_all_worlds_converged(self): + status = np.array( + [(True, 2, 1.0e-5), (False, 7, 2.0e-3)], + dtype=[("converged", np.bool_), ("iterations", np.int32), ("r_d", np.float32)], + ) + + reduced = _reduce_solver_status(status) + + self.assertFalse(reduced["converged"]) + self.assertEqual(reduced["iterations"], 7) + self.assertAlmostEqual(reduced["r_d"], 2.0e-3) + + def test_01_dvi_solve_dense_dual_problem(self): + builder = basics.build_boxes_fourbar() + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=0, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=None, + jacobians=jacobians, + ) + + dynamics_config = kamino_config.ConstrainedDynamicsConfig(preconditioning=True) + problem = DualProblem( + model=model, + data=data, + limits=limits, + contacts=detector.contacts, + jacobians=jacobians, + config=DualProblem.Config(dynamics=dynamics_config), + solver=LLTBlockedSolver, + sparse=False, + ) + problem.build(model=model, data=data, limits=limits, contacts=detector.contacts, jacobians=jacobians) + + solver = DVISolver( + model=model, + config=kamino_config.DVISolverConfig(max_iterations=1000, tolerance=1e-5, omega=1.0), + warmstart=WarmStartMode.NONE, + collect_info=True, + ) + solver.reset() + scratch = solver.data.state + for array in ( + scratch.v_aug, + scratch.s, + scratch.scratch, + scratch.bilateral_rhs, + scratch.bilateral_solution, + scratch.bilateral_preconditioner, + ): + array.fill_(float("nan")) + scratch.bilateral_active_dim.fill_(-1) + scratch.contact_colors.fill_(-1) + scratch.contact_num_colors.fill_(-1) + solver.coldstart() + solver.solve(problem) + _check_solution_matches_dual_problem(self, problem, solver) + np.testing.assert_array_equal(solver.data.info.status.numpy(), solver.data.status.numpy()) + + def test_02_public_solver_step_with_dvi(self): + builder = newton.ModelBuilder() + SolverKamino.register_custom_attributes(builder) + builder.default_shape_cfg.margin = 0.0 + builder.default_shape_cfg.gap = 0.0 + builder.begin_world() + body = builder.add_link( + label="link", + mass=1.0, + xform=wp.transformf(wp.vec3f(0.0, 0.0, 1.0), wp.quat_identity(dtype=wp.float32)), + ) + builder.add_shape_box(label="box", body=body, hx=0.1, hy=0.1, hz=0.1) + joint = builder.add_joint_revolute( + label="hinge", + parent=-1, + child=body, + axis=newton.Axis.Y, + parent_xform=wp.transformf(wp.vec3f(0.0, 0.0, 1.0), wp.quat_identity(dtype=wp.float32)), + child_xform=wp.transformf(wp.vec3f(0.0, 0.0, 0.0), wp.quat_identity(dtype=wp.float32)), + ) + builder.add_articulation([joint]) + builder.end_world() + model = builder.finalize(device=self.device) + + config = SolverKamino.Config( + dynamics_solver="dvi", + dvi=kamino_config.DVISolverConfig(max_iterations=500, tolerance=1e-5), + collect_solver_info=True, + ) + solver = SolverKamino(model, config=config) + state_in = model.state() + state_out = model.state() + solver.step(state_in, state_out, control=None, contacts=None, dt=1e-3) + body_q = state_out.body_q.numpy() + body_qd = state_out.body_qd.numpy() + self.assertTrue(np.all(np.isfinite(body_q))) + self.assertTrue(np.all(np.isfinite(body_qd))) + self.assertIsInstance(solver._solver_kamino.solver_fd, DVISolver) + self.assertFalse(solver._solver_kamino.config.dynamics.preconditioning) + self.assertIsNotNone(solver._solver_kamino.solver_fd.data.info) + + def test_03_dvi_solve_single_contact(self): + builder = basics.build_box_on_plane() + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=1, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + self.assertGreater(int(detector.contacts.model_active_contacts.numpy()[0]), 0) + + problem = DualProblem( + model=model, + data=data, + limits=limits, + contacts=detector.contacts, + jacobians=jacobians, + solver=LLTBlockedSolver, + sparse=False, + ) + problem.build(model=model, data=data, limits=limits, contacts=detector.contacts, jacobians=jacobians) + + solver = DVISolver( + model=model, + config=kamino_config.DVISolverConfig(max_iterations=200, tolerance=1e-4), + warmstart=WarmStartMode.NONE, + ) + solver.reset() + solver.coldstart() + solver.solve(problem) + status = solver.data.status.numpy()[0] + self.assertEqual(int(status["converged"]), 1) + self.assertLessEqual(int(status["iterations"]), _status_iteration_budget(solver, 0)) + self.assertTrue(np.all(np.isfinite(solver.data.solution.lambdas.numpy()))) + self.assertTrue(np.all(np.isfinite(solver.data.solution.v_plus.numpy()))) + + def test_03b_dvi_contact_block_preconditioner_smoke(self): + builder = basics.build_boxes_hinged() + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=8, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + self.assertGreater(int(detector.contacts.model_active_contacts.numpy()[0]), 0) + + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + + def solve_with_block_preconditioner(use_colored_contacts: bool) -> DVISolver: + solver = DVISolver( + model=model, + config=kamino_config.DVISolverConfig( + max_iterations=200, + tolerance=1e-4, + contact_block_preconditioner=True, + ), + warmstart=WarmStartMode.NONE, + ) + solver.reset() + solver.coldstart() + if use_colored_contacts: + solver.set_contacts(detector.contacts) + solver.solve(problem) + return solver + + contact_paths = [False] + if self.device.is_cuda: + contact_paths.append(True) + for use_colored_contacts in contact_paths: + with self.subTest(use_colored_contacts=use_colored_contacts): + solver = solve_with_block_preconditioner(use_colored_contacts) + + if use_colored_contacts: + self.assertGreater(int(solver.data.state.contact_num_colors.numpy()[0]), 0) + else: + self.assertEqual(int(solver.data.state.contact_num_colors.numpy()[0]), 0) + _assert_solution_finite(self, solver) + _check_solution_matches_dual_problem(self, problem, solver) + + def test_03c_dvi_noncolored_contact_jacobi_uses_configured_omega(self): + builder = basics.build_boxes_hinged() + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=8, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + self.assertGreater(int(detector.contacts.model_active_contacts.numpy()[0]), 0) + + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + + def solve_normal_lambda(contact_jacobi_omega: float) -> float: + solver = DVISolver( + model=model, + config=kamino_config.DVISolverConfig( + tolerance=0.0, + regularization=1e-5, + block_iterations=1, + contact_iterations=1, + contact_jacobi_omega=contact_jacobi_omega, + contact_jacobi_relaxation=1.0, + ), + warmstart=WarmStartMode.NONE, + ) + solver.reset() + solver.coldstart() + solver.solve(problem) + self.assertEqual(int(solver.data.state.contact_num_colors.numpy()[0]), 0) + lambdas = extract_problem_vector( + problem.delassus, solver.data.solution.lambdas.numpy(), only_active_dims=True + )[0] + contact_start = int(problem.data.ccgo.numpy()[0]) + contact_count = int(problem.data.nc.numpy()[0]) + normal_lambdas = lambdas[contact_start + 2 : contact_start + 3 * contact_count : 3] + return float(np.sum(normal_lambdas)) + + lambda_slow = solve_normal_lambda(0.1) + lambda_fast = solve_normal_lambda(0.8) + + self.assertTrue(np.isfinite(lambda_slow)) + self.assertTrue(np.isfinite(lambda_fast)) + self.assertGreater(lambda_fast, lambda_slow) + + def test_03d_dvi_direct_block_honors_per_world_iteration_counts(self): + builder = builder_utils.make_homogeneous_builder( + num_worlds=3, + build_fn=basics.build_boxes_hinged, + ) + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=8, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + self.assertTrue(np.all(detector.contacts.world_active_contacts.numpy() > 0)) + + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + configs = [ + kamino_config.DVISolverConfig( + tolerance=0.0, + regularization=1e-5, + block_iterations=1, + contact_iterations=1, + contact_jacobi_relaxation=1.0, + ), + kamino_config.DVISolverConfig( + tolerance=0.0, + regularization=1e-5, + block_iterations=3, + contact_iterations=1, + contact_jacobi_relaxation=1.0, + ), + kamino_config.DVISolverConfig( + tolerance=0.0, + regularization=1e-5, + block_iterations=1, + contact_iterations=3, + contact_jacobi_relaxation=1.0, + ), + ] + + def solve_normal_sums(use_colored_contacts: bool) -> list[float]: + solver = DVISolver(model=model, config=configs, warmstart=WarmStartMode.NONE) + solver.reset() + solver.coldstart() + if use_colored_contacts: + solver.set_contacts(detector.contacts) + solver.solve(problem) + status = solver.data.status.numpy() + self.assertEqual([int(status[wid]["iterations"]) for wid in range(3)], [1, 3, 3]) + if use_colored_contacts: + self.assertTrue(np.all(solver.data.state.contact_num_colors.numpy() > 0)) + else: + self.assertTrue(np.all(solver.data.state.contact_num_colors.numpy() == 0)) + np.testing.assert_array_equal( + solver.data.state.bilateral_active_dim.numpy(), + problem.data.njc.numpy(), + ) + + lambdas = extract_problem_vector( + problem.delassus, solver.data.solution.lambdas.numpy(), only_active_dims=True + ) + ccgo = problem.data.ccgo.numpy().astype(int) + nc = problem.data.nc.numpy().astype(int) + return [float(np.sum(lambdas[wid][ccgo[wid] + 2 : ccgo[wid] + 3 * nc[wid] : 3])) for wid in range(3)] + + contact_paths = [False] + if self.device.is_cuda: + contact_paths.append(True) + for use_colored_contacts in contact_paths: + with self.subTest(use_colored_contacts=use_colored_contacts): + normal_sums = solve_normal_sums(use_colored_contacts) + + self.assertGreater(normal_sums[1], normal_sums[0]) + self.assertGreater(normal_sums[2], normal_sums[0]) + + def test_03d2_dvi_direct_block_finishes_with_bilateral_solve(self): + builder = basics.build_boxes_hinged() + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=8, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + self.assertGreater(int(detector.contacts.world_active_contacts.numpy()[0]), 0) + + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + solver = DVISolver( + model=model, + config=kamino_config.DVISolverConfig( + tolerance=0.0, + regularization=1e-5, + block_iterations=1, + contact_iterations=1, + contact_jacobi_relaxation=1.0, + ), + warmstart=WarmStartMode.NONE, + ) + solver.reset() + solver.coldstart() + solver.solve(problem) + v_plus = extract_problem_vector(problem.delassus, solver.data.solution.v_plus.numpy(), only_active_dims=True)[0] + njc = int(problem.data.njc.numpy()[0]) + status = solver.data.status.numpy()[0] + + self.assertGreater(njc, 0) + self.assertLess(float(np.max(np.abs(v_plus[:njc]))), 1e-6) + self.assertLess(float(status["r_b"]), 1e-6) + + def test_03e_dvi_direct_block_no_unilateral_rows_reports_single_iteration(self): + builder = basics.build_box_pendulum(ground=False) + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=4, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + self.assertGreater(int(model.info.num_joint_cts.numpy()[0]), 0) + self.assertEqual(int(problem.data.nl.numpy()[0]), 0) + self.assertEqual(int(problem.data.nc.numpy()[0]), 0) + + solver = DVISolver( + model=model, + config=kamino_config.DVISolverConfig( + tolerance=1e-4, + regularization=1e-5, + block_iterations=7, + contact_iterations=3, + ), + warmstart=WarmStartMode.NONE, + ) + solver.reset() + solver.coldstart() + solver.solve(problem) + status = solver.data.status.numpy()[0] + self.assertEqual(int(status["converged"]), 1) + self.assertEqual(int(status["iterations"]), 1) + self.assertEqual(int(solver.data.state.bilateral_active_dim.numpy()[0]), 0) + _check_solution_matches_dual_problem(self, problem, solver) + + def test_03f_dvi_bilateral_only_solve_resets_stale_status(self): + builder = basics.build_box_pendulum(ground=False) + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=0, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=None, + jacobians=jacobians, + ) + + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + self.assertGreater(int(model.info.num_joint_cts.numpy()[0]), 0) + self.assertEqual(int(problem.data.nl.numpy()[0]), 0) + self.assertEqual(int(problem.data.nc.numpy()[0]), 0) + + solver = DVISolver( + model=model, + config=kamino_config.DVISolverConfig( + tolerance=1e-4, + regularization=1e-5, + contact_iterations=5, + ), + warmstart=WarmStartMode.NONE, + ) + solver.reset() + solver.coldstart() + wp.launch( + kernel=_initialize_dvi_status, + dim=solver.size.num_worlds, + inputs=[ + solver.data.config, + solver.data.status, + ], + device=self.device, + ) + self.assertEqual(int(solver.data.status.numpy()[0]["iterations"]), 5) + + solver.solve(problem) + status = solver.data.status.numpy()[0] + self.assertEqual(int(status["converged"]), 1) + self.assertEqual(int(status["iterations"]), 1) + _check_solution_matches_dual_problem(self, problem, solver) + + def test_03g_dvi_contact_coloring_separates_dynamic_conflicts(self): + problem_nc = wp.array([5], dtype=wp.int32, device=self.device) + problem_cio = wp.array([0], dtype=wp.int32, device=self.device) + contact_bid_ab = wp.array( + [ + wp.vec2i(0, -1), + wp.vec2i(0, 1), + wp.vec2i(2, -1), + wp.vec2i(-1, -1), + wp.vec2i(1, -1), + ], + dtype=wp.vec2i, + device=self.device, + ) + contact_colors = wp.full(shape=5, value=-1, dtype=wp.int32, device=self.device) + contact_num_colors = wp.zeros(shape=1, dtype=wp.int32, device=self.device) + + wp.launch( + kernel=_color_dvi_contacts, + dim=1, + inputs=[ + problem_nc, + problem_cio, + contact_bid_ab, + contact_colors, + contact_num_colors, + ], + device=self.device, + ) + colors = contact_colors.numpy() + num_colors = int(contact_num_colors.numpy()[0]) + self.assertGreaterEqual(num_colors, 2) + self.assertTrue(np.all(colors >= 0)) + self.assertNotEqual(colors[0], colors[1]) + self.assertNotEqual(colors[1], colors[4]) + self.assertLess(colors[3], num_colors) + + def test_03i_dvi_coldstart_is_repeatable(self): + for sparse in (False, True): + with self.subTest(sparse=sparse): + test = TestSetup( + builder_fn=basics.build_boxes_hinged, + max_world_contacts=8, + gravity=True, + perturb=True, + device=self.device, + sparse=sparse, + ) + test.build() + config = SolverKamino.Config( + dynamics_solver="dvi", + sparse_dynamics=sparse, + sparse_jacobian=sparse, + ).dvi + solver = _solve_dvi(test.model, test.problem, config=config, setup=test) + first_lambdas = solver.data.solution.lambdas.numpy().copy() + first_v_plus = solver.data.solution.v_plus.numpy().copy() + first_status = solver.data.status.numpy().copy() + + test.build() + solver.reset() + solver.coldstart() + solver.solve(test.problem) + + np.testing.assert_allclose(solver.data.solution.lambdas.numpy(), first_lambdas, rtol=0.0, atol=1e-6) + np.testing.assert_allclose(solver.data.solution.v_plus.numpy(), first_v_plus, rtol=0.0, atol=1e-6) + status = solver.data.status.numpy() + np.testing.assert_array_equal(status["converged"], first_status["converged"]) + np.testing.assert_array_equal(status["iterations"], first_status["iterations"]) + for residual in ("r_p", "r_d", "r_c", "r_b"): + np.testing.assert_allclose(status[residual], first_status[residual], rtol=1e-5, atol=1e-8) + + def test_04_dvi_solve_active_joint_limit(self): + builder = testing.build_unary_revolute_joint_test(limits=True, ground=False) + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=0, + sparse=False, + ) + update_containers(model=model, data=data, state=state, limits=limits, detector=None, jacobians=None) + + q_j = data.joints.q_j.numpy() + q_j[:] = 1.0 + data.joints.q_j.assign(q_j) + limits.detect(q_j=data.joints.q_j) + update_constraints_info(model=model, data=data) + jacobians.build(model=model, data=data, limits=limits.data, contacts=None) + self.assertGreater(int(limits.model_active_limits.numpy()[0]), 0) + + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + solver = _solve_dvi( + model, + problem, + config=kamino_config.DVISolverConfig(max_iterations=8, tolerance=1e-4, regularization=1e-5), + ) + + _assert_solver_status_converged(self, solver) + iterations = int(solver.data.status.numpy()[0]["iterations"]) + self.assertEqual(iterations, solver.config[0].block_iterations * solver.config[0].contact_iterations) + self.assertGreater(iterations, solver.config[0].max_iterations) + _check_solution_matches_dual_problem(self, problem, solver) + + def test_07_dvi_singular_limit_rows_remain_finite(self): + model, data, _state, limits, contacts = make_test_problem_fourbar( + device=self.device, + max_world_contacts=0, + with_limits=True, + with_contacts=False, + ) + jacobians = DenseSystemJacobians(model=model, limits=limits, contacts=contacts) + jacobians.build(model=model, data=data, limits=limits.data, contacts=None) + self.assertGreater(int(limits.model_active_limits.numpy()[0]), 0) + + problem = _make_dense_dual_problem(model, data, limits, contacts, jacobians) + solver = _solve_dvi(model, problem) + + status = solver.data.status.numpy()[0] + self.assertEqual(int(status["converged"]), 0) + _assert_solution_finite(self, solver) + lambdas_np = extract_problem_vector( + problem.delassus, solver.data.solution.lambdas.numpy(), only_active_dims=True + )[0] + limit_start = int(problem.data.lcgo.numpy()[0]) + limit_count = int(problem.data.nl.numpy()[0]) + limit_lambdas = lambdas_np[limit_start : limit_start + limit_count] + self.assertLess(float(np.max(np.abs(limit_lambdas))), 1.0) + + def test_08_public_solver_short_rollout_with_dvi(self): + builder = newton.ModelBuilder() + SolverKamino.register_custom_attributes(builder) + builder.default_shape_cfg.margin = 0.0 + builder.default_shape_cfg.gap = 0.0 + builder.begin_world() + body = builder.add_link( + label="link", + mass=1.0, + xform=wp.transformf(wp.vec3f(0.0, 0.0, 1.0), wp.quat_identity(dtype=wp.float32)), + ) + builder.add_shape_box(label="box", body=body, hx=0.1, hy=0.1, hz=0.1) + joint = builder.add_joint_revolute( + label="hinge", + parent=-1, + child=body, + axis=newton.Axis.Y, + parent_xform=wp.transformf(wp.vec3f(0.0, 0.0, 1.0), wp.quat_identity(dtype=wp.float32)), + child_xform=wp.transformf(wp.vec3f(0.0, 0.0, 0.0), wp.quat_identity(dtype=wp.float32)), + ) + builder.add_articulation([joint]) + builder.end_world() + model = builder.finalize(device=self.device) + + config = SolverKamino.Config( + dynamics_solver="dvi", + dvi=kamino_config.DVISolverConfig(max_iterations=300, tolerance=1e-4), + ) + solver = SolverKamino(model, config=config) + state_in = model.state() + state_out = model.state() + for _ in range(8): + solver.step(state_in, state_out, control=None, contacts=None, dt=1e-3) + state_in, state_out = state_out, state_in + self.assertTrue(np.all(np.isfinite(state_in.body_q.numpy()))) + self.assertTrue(np.all(np.isfinite(state_in.body_qd.numpy()))) + self.assertIsInstance(solver._solver_kamino.solver_fd, DVISolver) + + def test_08a_public_solver_heterogeneous_contact_rollout_with_dvi(self): + builder = newton.ModelBuilder() + SolverKamino.register_custom_attributes(builder) + public_basics.make_basics_heterogeneous_builder(builder=builder, ground=True) + model = builder.finalize(device=self.device, skip_validation_joints=True) + + config = SolverKamino.Config( + dynamics_solver="dvi", + use_collision_detector=True, + collision_detector=kamino_config.CollisionDetectorConfig( + max_contacts=64 * model.world_count, + max_contacts_per_world=64, + max_contacts_per_pair=16, + ), + ) + solver = SolverKamino(model, config=config) + state_in = model.state() + state_out = model.state() + control = model.control() + + contact_seen = False + for _ in range(24): + solver.step(state_in, state_out, control=control, contacts=None, dt=1.0e-3) + state_in, state_out = state_out, state_in + contact_seen = contact_seen or bool(np.any(solver._contacts_kamino.world_active_contacts.numpy() > 0)) + + status = solver._solver_kamino.solver_fd.data.status.numpy() + self.assertTrue(contact_seen) + self.assertTrue(np.all(np.isfinite(state_in.body_q.numpy()))) + self.assertTrue(np.all(np.isfinite(state_in.body_qd.numpy()))) + self.assertTrue(np.all(np.isfinite(status["r_p"]))) + self.assertTrue(np.all(np.isfinite(status["r_d"]))) + self.assertLess(float(np.max(np.abs(state_in.body_qd.numpy()))), 100.0) + self.assertIsInstance(solver._solver_kamino.solver_fd, DVISolver) + self.assertFalse(config.sparse_dynamics) + self.assertTrue(config.sparse_jacobian) + self.assertEqual(config.dynamics.linear_solver_type, "LLTBRCM") + + def test_03a_sparse_dvi_filtered_matvec_matches_full_rows(self): + builder = basics.build_box_on_plane() + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=4, + sparse=True, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + self.assertGreater(int(detector.contacts.model_active_contacts.numpy()[0]), 0) + + problem = _make_sparse_dual_problem(model, data, limits, detector.contacts, jacobians) + solver = DVISolver( + model=model, + data=data, + limits=limits, + contacts=detector.contacts, + jacobians=jacobians, + config=kamino_config.DVISolverConfig( + tolerance=0.0, + regularization=1e-5, + block_iterations=1, + contact_iterations=1, + ), + warmstart=WarmStartMode.NONE, + ) + solver.reset() + + lambdas = np.linspace(-0.25, 0.5, problem.data.v_f.shape[0], dtype=np.float32) + solver.data.solution.lambdas.assign(lambdas) + + full = wp.zeros_like(problem.data.v_f) + problem.delassus.matvec(solver.data.solution.lambdas, full, solver.all_worlds_mask) + full_np = full.numpy() + + _sparse_delassus_matvec_rows(solver, problem, _SPARSE_DELASSUS_ROWS_JOINTS) + joint_np = solver.data.state.v_aug.numpy() + _sparse_delassus_matvec_rows(solver, problem, _SPARSE_DELASSUS_ROWS_UNILATERAL) + unilateral_np = solver.data.state.v_aug.numpy() + + dim = int(problem.data.dim.numpy()[0]) + njc = int(problem.data.njc.numpy()[0]) + np.testing.assert_allclose(joint_np[:njc], full_np[:njc], rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(unilateral_np[njc:dim], full_np[njc:dim], rtol=1e-5, atol=1e-5) + + def test_05_dvi_solve_multi_world_contacts(self): + builder = builder_utils.make_homogeneous_builder( + num_worlds=4, + build_fn=basics.build_box_on_plane, + ground=True, + ) + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=4, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + self.assertTrue(np.all(detector.contacts.world_active_contacts.numpy() > 0)) + + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + solver = _solve_dvi(model, problem) + + _assert_solver_status_converged(self, solver) + _check_solution_matches_dual_problem(self, problem, solver) + + def test_06_dvi_warmstart_modes(self): + builder = basics.build_box_on_plane() + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=4, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + + internal_solver = _solve_dvi(model, problem, warmstart=WarmStartMode.INTERNAL) + cold_iterations = int(internal_solver.data.status.numpy()[0]["iterations"]) + _assert_solver_status_converged(self, internal_solver) + + problem.build(model=model, data=data, limits=limits, contacts=detector.contacts, jacobians=jacobians) + internal_solver.warmstart(problem, model, data) + internal_solver.solve(problem) + _assert_solver_status_converged(self, internal_solver) + self.assertLessEqual(int(internal_solver.data.status.numpy()[0]["iterations"]), cold_iterations) + + unpack_constraint_solutions( + lambdas=internal_solver.data.solution.lambdas, + v_plus=internal_solver.data.solution.v_plus, + model=model, + data=data, + limits=limits, + contacts=detector.contacts, + ) + container_solver = DVISolver( + model=model, + config=kamino_config.DVISolverConfig(max_iterations=300, tolerance=1e-4, regularization=1e-5), + warmstart=WarmStartMode.CONTAINERS, + ) + problem.build(model=model, data=data, limits=limits, contacts=detector.contacts, jacobians=jacobians) + container_solver.warmstart(problem, model, data, limits, detector.contacts) + container_solver.solve(problem) + _assert_solver_status_converged(self, container_solver) + _check_solution_matches_dual_problem(self, problem, container_solver) + + def test_06a_dvi_masked_reset_preserves_unselected_worlds(self): + builder = builder_utils.make_homogeneous_builder( + num_worlds=3, + build_fn=basics.build_box_on_plane, + ground=True, + ) + model, data, state, limits, detector, jacobians = make_containers( + builder=builder, + device=self.device, + max_world_contacts=4, + sparse=False, + ) + update_containers( + model=model, + data=data, + state=state, + limits=limits, + detector=detector, + jacobians=jacobians, + ) + problem = _make_dense_dual_problem(model, data, limits, detector.contacts, jacobians) + solver = _solve_dvi(model, problem) + lambdas_before = solver.data.solution.lambdas.numpy().copy() + v_plus_before = solver.data.solution.v_plus.numpy().copy() + + world_mask = wp.array([False, True, False], dtype=wp.bool, device=self.device) + solver.reset(problem=problem, world_mask=world_mask) + + lambdas_after = extract_problem_vector( + problem.delassus, solver.data.solution.lambdas.numpy(), only_active_dims=False + ) + v_plus_after = extract_problem_vector( + problem.delassus, solver.data.solution.v_plus.numpy(), only_active_dims=False + ) + lambdas_before = extract_problem_vector(problem.delassus, lambdas_before, only_active_dims=False) + v_plus_before = extract_problem_vector(problem.delassus, v_plus_before, only_active_dims=False) + np.testing.assert_array_equal(lambdas_after[0], lambdas_before[0]) + np.testing.assert_array_equal(lambdas_after[2], lambdas_before[2]) + np.testing.assert_array_equal(v_plus_after[0], v_plus_before[0]) + np.testing.assert_array_equal(v_plus_after[2], v_plus_before[2]) + np.testing.assert_array_equal(lambdas_after[1], np.zeros_like(lambdas_after[1])) + np.testing.assert_array_equal(v_plus_after[1], np.zeros_like(v_plus_after[1])) + + def test_12_dvi_opening_contact_releases_warmstarted_force(self): + if not self.device.is_cuda: + self.skipTest("DVI colored contact release regression uses the CUDA graph-colored path") + + radius = 0.1 + separation = 0.005 + gap = 0.03 + z = radius + separation + + builder = newton.ModelBuilder(up_axis=newton.Axis.Z) + SolverKamino.register_custom_attributes(builder) + shape_cfg = newton.ModelBuilder.ShapeConfig(gap=gap, margin=0.0) + body = builder.add_link( + xform=wp.transform(p=wp.vec3(0.0, 0.0, z), q=wp.quat_identity()), + mass=1.0, + ) + builder.add_shape_sphere(body=body, radius=radius, cfg=shape_cfg) + joint = builder.add_joint_prismatic( + parent=-1, + child=body, + axis=newton.Axis.Z, + parent_xform=wp.transform(p=wp.vec3(0.0, 0.0, z), q=wp.quat_identity()), + child_xform=wp.transform_identity(), + limit_lower=-10.0, + limit_upper=10.0, + ) + builder.add_articulation([joint]) + builder.add_ground_plane(cfg=shape_cfg) + model = builder.finalize(device=self.device) + + joint_qd = model.joint_qd.numpy() + joint_qd[:] = 1.0 + model.joint_qd.assign(joint_qd) + + state_0 = model.state() + state_1 = model.state() + newton.eval_fk(model, model.joint_q, model.joint_qd, state_0) + + config = SolverKamino.Config( + use_collision_detector=True, + collision_detector=kamino_config.CollisionDetectorConfig( + max_contacts_per_world=8, + max_contacts_per_pair=8, + default_gap=gap, + ), + dynamics_solver="dvi", + dvi=kamino_config.DVISolverConfig( + max_iterations=300, + tolerance=1e-5, + regularization=1e-5, + block_iterations=32, + contact_iterations=4, + ), + ) + solver = SolverKamino(model, config=config) + + solver.step(state_0, state_1, control=None, contacts=None, dt=1e-3) + self.assertEqual(int(solver._contacts_kamino.model_active_contacts.numpy()[0]), 1) + + cache = solver._solver_kamino._ws_contacts.cache + self.assertIsNotNone(cache) + reaction = cache.reaction.numpy() + reaction[:, :] = 0.0 + reaction[0, 2] = 10000.0 + cache.reaction.assign(reaction) + velocity = cache.velocity.numpy() + velocity[:, :] = 0.0 + cache.velocity.assign(velocity) + + solver.step(state_1, state_0, control=None, contacts=None, dt=1e-3) + + contact_count = int(solver._contacts_kamino.model_active_contacts.numpy()[0]) + gaps = solver._contacts_kamino.gapfunc.numpy()[:contact_count, 3] + contact_velocity = solver._contacts_kamino.velocity.numpy()[:contact_count, 2] + contact_reaction = solver._contacts_kamino.reaction.numpy()[:contact_count, 2] + opening = (gaps > 0.0) & (contact_velocity > 0.0) + + self.assertTrue(np.any(opening)) + self.assertLess(float(np.max(np.abs(contact_reaction[opening]))), 1e-3) + self.assertLess(float(abs(state_0.body_qd.numpy()[0, 2])), 2.0) + self.assertEqual(int(solver._solver_kamino.solver_fd.data.status.numpy()[0]["converged"]), 1) + + def test_03h_dvi_canonical_contact_solution_metrics(self): + for builder_fn, max_world_contacts in ( + (basics.build_box_on_plane, 4), + (basics.build_boxes_hinged, 8), + ): + for sparse in (False, True): + with self.subTest(builder=builder_fn.__name__, sparse=sparse): + test = TestSetup( + builder_fn=builder_fn, + max_world_contacts=max_world_contacts, + gravity=True, + perturb=True, + device=self.device, + sparse=sparse, + ) + test.build() + config = SolverKamino.Config( + dynamics_solver="dvi", + sparse_dynamics=sparse, + sparse_jacobian=sparse, + ).dvi + solver = _solve_dvi(test.model, test.problem, config=config, setup=test) + solution_metrics = _evaluate_solution_metrics(test, solver) + + _assert_solution_finite(self, solver) + for name, value in solution_metrics.items(): + self.assertTrue(np.isfinite(value), msg=f"{name}={value}") + + # DVI trades some contact accuracy for throughput, but its + # solution must still satisfy dynamics and cone feasibility. + self.assertLess(solution_metrics["r_eom"], 1.0e-4, msg=str(solution_metrics)) + self.assertLess(solution_metrics["r_kinematics"], 1.0e-4, msg=str(solution_metrics)) + self.assertLess(solution_metrics["r_cts_joints"], 1.0e-4, msg=str(solution_metrics)) + self.assertLess(solution_metrics["r_v_plus"], 1.0e-4, msg=str(solution_metrics)) + self.assertLess(solution_metrics["r_ncp_primal"], 1.0e-4, msg=str(solution_metrics)) + self.assertLess(solution_metrics["r_ncp_dual"], 1.0e-2, msg=str(solution_metrics)) + self.assertLess(solution_metrics["r_ncp_compl"], 1.0e-2, msg=str(solution_metrics)) + self.assertLess(solution_metrics["r_vi_natmap"], 1.0e-2, msg=str(solution_metrics)) + + def test_08b_dr_legs_contact_capacity_scales_with_world_count(self): + if not self.device.is_cuda: + self.skipTest("Dr Legs multi-world capacity regression uses the CUDA graph path") + + from types import SimpleNamespace # noqa: PLC0415 + + from newton.examples.kamino.example_kamino_robot_dr_legs import Example # noqa: PLC0415 + from newton.viewer import ViewerNull # noqa: PLC0415 + + world_count = 3 + args = SimpleNamespace( + world_count=world_count, + use_kamino_contacts=False, + dynamics_solver="dvi", + dvi_contact_block_preconditioner=False, + dvi_contact_jacobi_omega=0.45, + dvi_contact_jacobi_relaxation=0.9, + ) + example = Example(ViewerNull(num_frames=1), args) + + expected_capacity = 72 * world_count + self.assertEqual(example.model.rigid_contact_max, expected_capacity) + self.assertEqual(example.contacts.rigid_contact_max, expected_capacity) + self.assertEqual(example.collision_pipeline.rigid_contact_max, expected_capacity) + + def test_09_dr_legs_dvi_first_contact_remains_finite(self): + if not self.device.is_cuda: + self.skipTest("Dr Legs DVI first-contact regression uses the CUDA graph path") + + from types import SimpleNamespace # noqa: PLC0415 + + from newton.examples.kamino.example_kamino_robot_dr_legs import Example # noqa: PLC0415 + from newton.viewer import ViewerNull # noqa: PLC0415 + + args = SimpleNamespace( + world_count=1, + use_kamino_contacts=True, + dynamics_solver="dvi", + dvi_contact_block_preconditioner=False, + dvi_contact_jacobi_omega=0.45, + dvi_contact_jacobi_relaxation=0.9, + ) + example = Example(ViewerNull(num_frames=1), args) + + contact_seen = False + color_checked = False + for _ in range(12): + example.step() + body_q = example.state_0.body_q.numpy() + body_qd = example.state_0.body_qd.numpy() + lambdas = example.solver._solver_kamino.solver_fd.data.solution.lambdas.numpy() + kamino_contacts = example.solver._contacts_kamino + contact_count = int(kamino_contacts.world_active_contacts.numpy()[0]) + contact_seen = contact_seen or contact_count > 0 + if contact_count > 0 and not example.config.sparse_dynamics: + solver_fd = example.solver._solver_kamino.solver_fd + color_count = int(solver_fd.data.state.contact_num_colors.numpy()[0]) + colors = solver_fd.data.state.contact_colors.numpy() + bid_ab = kamino_contacts.bid_AB.numpy() + self.assertGreater(color_count, 0) + self.assertTrue(np.all(colors[:contact_count] >= 0)) + for ci in range(contact_count): + bodies_i = {int(bid_ab[ci][0]), int(bid_ab[ci][1])} - {-1} + for cj in range(ci): + if colors[ci] == colors[cj]: + bodies_j = {int(bid_ab[cj][0]), int(bid_ab[cj][1])} - {-1} + self.assertFalse(bodies_i & bodies_j) + color_checked = True + elif contact_count > 0: + color_checked = True + + self.assertTrue(np.all(np.isfinite(body_q))) + self.assertTrue(np.all(np.isfinite(body_qd))) + self.assertTrue(np.all(np.isfinite(lambdas))) + self.assertLess(float(np.max(np.abs(body_qd))), 100.0) + self.assertLess(float(np.max(np.abs(lambdas))), 100.0) + + self.assertTrue(contact_seen) + self.assertTrue(color_checked) + + def test_10_dr_legs_dvi_tipped_contact_does_not_creep(self): + if not self.device.is_cuda: + self.skipTest("Dr Legs DVI tipped-contact regression uses the CUDA graph path") + + from types import SimpleNamespace # noqa: PLC0415 + + from newton.examples.kamino.example_kamino_robot_dr_legs import Example # noqa: PLC0415 + from newton.viewer import ViewerNull # noqa: PLC0415 + + args = SimpleNamespace( + world_count=1, + use_kamino_contacts=True, + dynamics_solver="dvi", + dvi_contact_block_preconditioner=False, + dvi_contact_jacobi_omega=0.45, + dvi_contact_jacobi_relaxation=0.9, + ) + example = Example(ViewerNull(num_frames=1), args) + + q_tip = wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), float(np.pi * 0.5)) + example.base_q.assign([wp.transformf((0.0, 0.0, 0.25), q_tip)]) + reset_config = SolverKamino.ResetConfig(base_pose=SolverKamino.ResetConfig.FromBaseQ(example.base_q)) + example.solver.reset(state=example.state_0, config=reset_config) + example.solver.reset(state=example.state_1, config=reset_config) + example.capture() + + base_start = example.state_0.body_q.numpy()[0, :3].copy() + contact_seen = False + post_settle_penetration = [] + post_settle_xy = [] + for step_idx in range(400): + example.step() + contact_seen = contact_seen or int(example.contacts.rigid_contact_count.numpy()[0]) > 0 + contacts_kamino = example.solver._contacts_kamino + contact_count = int(contacts_kamino.world_active_contacts.numpy()[0]) + if step_idx >= 40 and contact_count > 0: + gaps = contacts_kamino.gapfunc.numpy()[:contact_count, 3] + post_settle_penetration.append(float(max(0.0, -np.min(gaps)))) + if step_idx >= 200: + post_settle_xy.append(example.state_0.body_q.numpy()[0, :2].copy()) + + body_q = example.state_0.body_q.numpy() + body_qd = example.state_0.body_qd.numpy() + base_delta_xy = body_q[0, :2] - base_start[:2] + + self.assertTrue(contact_seen) + self.assertTrue(np.all(np.isfinite(body_q))) + self.assertTrue(np.all(np.isfinite(body_qd))) + self.assertLess(float(np.linalg.norm(base_delta_xy)), 0.008) + self.assertGreater(len(post_settle_penetration), 0) + self.assertLess(float(np.percentile(post_settle_penetration, 95)), 0.0035) + self.assertLess(float(np.linalg.norm(post_settle_xy[-1] - post_settle_xy[0])), 5.0e-4) + + def test_11_dr_legs_dvi_contact_force_balances_weight(self): + if not self.device.is_cuda: + self.skipTest("Dr Legs DVI contact-force regression uses the CUDA graph path") + + from types import SimpleNamespace # noqa: PLC0415 + + from newton._src.solvers.kamino._src.geometry.aggregation import ContactAggregation # noqa: PLC0415 + from newton.examples.kamino.example_kamino_robot_dr_legs import Example # noqa: PLC0415 + from newton.viewer import ViewerNull # noqa: PLC0415 + + args = SimpleNamespace( + world_count=1, + use_kamino_contacts=True, + dynamics_solver="dvi", + dvi_contact_block_preconditioner=False, + dvi_contact_jacobi_omega=0.45, + dvi_contact_jacobi_relaxation=0.9, + ) + example = Example(ViewerNull(num_frames=1), args) + + base_z = [] + for _ in range(180): + example.step() + base_z.append(float(example.state_0.body_q.numpy()[0, 2])) + + contacts_kamino = example.solver._contacts_kamino + aggregation = ContactAggregation(model=example.solver._model_kamino, contacts=contacts_kamino) + aggregation.compute() + + contact_count = int(contacts_kamino.world_active_contacts.numpy()[0]) + total_contact_force = aggregation.body_net_force.numpy()[0].sum(axis=0) + weight = float(example.model.body_mass.numpy().sum() * 9.81) + force_ratio = float(total_contact_force[2] / weight) + + self.assertGreater(contact_count, 0) + self.assertTrue(np.all(np.isfinite(total_contact_force))) + self.assertGreater(force_ratio, 0.95) + self.assertLess(force_ratio, 1.05) + z = np.array(base_z[60:], dtype=np.float64) + x = np.arange(z.size, dtype=np.float64) + residual = z - np.polyval(np.polyfit(x, z, 1), x) + self.assertLess(float(np.max(residual) - np.min(residual)), 0.001) + + +if __name__ == "__main__": + unittest.main() diff --git a/newton/_src/solvers/kamino/tests/test_solvers_forward_kinematics.py b/newton/_src/solvers/kamino/tests/test_solvers_forward_kinematics.py index 83905d5366..1c2e4fee5a 100644 --- a/newton/_src/solvers/kamino/tests/test_solvers_forward_kinematics.py +++ b/newton/_src/solvers/kamino/tests/test_solvers_forward_kinematics.py @@ -130,6 +130,40 @@ def eval_constraints(bodies_q_stepped_np): self.assertTrue(success) +class SparseJacobianSingleJointCheckForwardKinematics(unittest.TestCase): + def setUp(self): + if not test_context.setup_done: + setup_tests(clear_cache=False) + self.default_device = wp.get_device(test_context.device) + + def tearDown(self): + self.default_device = None + + def test_sparse_jacobian_matches_dense_for_single_joint_examples(self): + """Match dense and sparse Jacobians for every single-joint fixture.""" + test_name = "Single-joint sparse Jacobian assembly check" + rng = np.random.default_rng(42) + + def test_function(model: ModelKamino): + """Compare the dense and sparse Jacobians for a random body state.""" + bodies_q_np = rng.uniform(-1.0, 1.0, 7 * model.size.sum_of_num_bodies).astype("float32") + bodies_q = wp.from_numpy(bodies_q_np, dtype=wp.transformf, device=model.device) + actuators_q = wp.zeros( + shape=model.size.sum_of_num_actuated_joint_coords, dtype=wp.float32, device=model.device + ) + solver = ForwardKinematicsSolver(model, config=ForwardKinematicsSolver.Config(use_sparsity=True)) + transforms = solver.eval_position_control_transformations(actuators_q, None) + + jac_dense_np = solver.eval_kinematic_constraints_jacobian(bodies_q, transforms).numpy() + solver.assemble_sparse_jacobian(bodies_q, transforms) + jac_sparse_np = solver.sparse_jacobian.numpy() + rows, cols = solver.sparse_jacobian.dims.numpy()[0] + return np.allclose(jac_dense_np[0, :rows, :cols], jac_sparse_np[0], atol=1e-6, rtol=0.0) + + success = run_test_single_joint_examples(test_function, test_name, device=self.default_device) + self.assertTrue(success) + + class WorldMaskInitializationForwardKinematics(unittest.TestCase): def setUp(self): if not test_context.setup_done: @@ -391,7 +425,7 @@ def test_mechanism_FK_random_poses(self): model, num_poses, rng, - use_graph=self.has_cuda, + use_graph=self.has_cuda and not wp.config.verify_cuda, verbose=self.verbose, reset_state=True, use_incremental_solve=True, @@ -424,7 +458,7 @@ def test_dr_legs_FK_random_poses(self): seed = int(hashlib.sha256(test_name.encode("utf8")).hexdigest(), 16) rng = np.random.default_rng(seed) - # Load the DR TestMech and DR Legs models from the `newton-assets` repository + # Load the DR Legs model from the `newton-assets` repository asset_path = newton.utils.download_asset("disneyresearch") asset_file = str(asset_path / "dr_legs" / "usd" / "dr_legs_with_boxes.usda") builder = USDImporter().import_from(asset_file) @@ -440,7 +474,7 @@ def test_dr_legs_FK_random_poses(self): rng, max_angle=np.radians(10.0), # Angles too far from the initial pose lead to singularities max_ang_vel=np.radians(30.0), - use_graph=self.has_cuda, + use_graph=self.has_cuda and not wp.config.verify_cuda, verbose=self.verbose, reset_state=True, tolerance=1e-6, @@ -491,7 +525,7 @@ def test_heterogenous_model_FK_random_poses(self): rng, max_angle=np.radians(10.0), # Angles too far from the initial pose lead to singularities max_ang_vel=np.radians(30.0), - use_graph=self.has_cuda, + use_graph=self.has_cuda and not wp.config.verify_cuda, verbose=self.verbose, reset_state=True, use_incremental_solve=True, @@ -518,6 +552,46 @@ def setUp(self): def tearDown(self): self.default_device = None + def test_axis_joint_frames_update_after_notify(self): + """Synthetic axis frames match a fresh solver after model changes.""" + model = create_four_bar_tie_rod().finalize(device=self.default_device, requires_grad=False) + config = ForwardKinematicsSolver.Config(add_axis_joints=True) + solver = ForwardKinematicsSolver(model, config) + axis_body = int(solver.fk_axis_body.numpy()[0]) + source_joint = int(solver.fk_axis_source_joint_0.numpy()[0]) + + body_q = model.bodies.q_i_0.numpy() + body_q[axis_body] = np.array( + wp.transformf( + wp.vec3f(*body_q[axis_body, :3]), + wp.quat_from_axis_angle(wp.vec3f(0.0, 1.0, 0.0), 0.3), + ) + ) + model.bodies.q_i_0.assign(body_q) + if model.joints.bid_B.numpy()[source_joint] == axis_body: + joint_anchor = model.joints.B_r_Bj.numpy() + joint_anchor[source_joint] += np.array([0.05, -0.02, 0.01], dtype=np.float32) + model.joints.B_r_Bj.assign(joint_anchor) + else: + joint_anchor = model.joints.F_r_Fj.numpy() + joint_anchor[source_joint] += np.array([0.05, -0.02, 0.01], dtype=np.float32) + model.joints.F_r_Fj.assign(joint_anchor) + + solver.notify_model_changed(newton.ModelFlags.JOINT_PROPERTIES | newton.ModelFlags.BODY_PROPERTIES) + reference = ForwardKinematicsSolver(model, ForwardKinematicsSolver.Config(add_axis_joints=True)) + axis_joints = solver.fk_axis_joint.numpy() + + np.testing.assert_allclose( + solver.joints_X_Bj.numpy()[axis_joints], + reference.joints_X_Bj.numpy()[axis_joints], + atol=1e-6, + ) + np.testing.assert_allclose( + solver.joints_X_Fj.numpy()[axis_joints], + reference.joints_X_Fj.numpy()[axis_joints], + atol=1e-6, + ) + def test_four_bar_tie_rod_model_FK_random_poses(self): # Initialize RNG test_name = "Four-bar with tie rod FK random poses check" @@ -535,7 +609,7 @@ def test_four_bar_tie_rod_model_FK_random_poses(self): model, num_poses, rng, - use_graph=self.has_cuda, + use_graph=self.has_cuda and not wp.config.verify_cuda, verbose=self.verbose, reset_state=True, use_incremental_solve=True, @@ -587,7 +661,7 @@ def test_all_joints_example_FK_random_poses(self): model, num_poses, rng, - use_graph=self.has_cuda, + use_graph=self.has_cuda and not wp.config.verify_cuda, verbose=self.verbose, reset_state=True, use_incremental_solve=True, @@ -632,7 +706,7 @@ def test_all_joints_example_asymmetric_frames_FK_random_poses(self): model, num_poses, rng, - use_graph=self.has_cuda, + use_graph=self.has_cuda and not wp.config.verify_cuda, verbose=self.verbose, reset_state=True, use_incremental_solve=True, @@ -686,7 +760,7 @@ def test_cartpole_FK_random_poses(self): model, num_poses, rng, - use_graph=self.has_cuda, + use_graph=self.has_cuda and not wp.config.verify_cuda, verbose=self.verbose, reset_state=True, use_incremental_solve=True, @@ -756,7 +830,7 @@ def test_heterogenous_model_FK_random_poses(self): for wd_id in range(model.size.num_worlds): rows, cols = int(dims[wd_id][0]), int(dims[wd_id][1]) residual = jac_dense_np[wd_id, :rows, :cols] - jac_sparse_np[wd_id] - self.assertTrue(np.max(np.abs(residual)) < 1e-10) + self.assertTrue(np.max(np.abs(residual)) < 1e-6) ### diff --git a/newton/_src/solvers/kamino/tests/test_solvers_metrics.py b/newton/_src/solvers/kamino/tests/test_solvers_metrics.py index b71c71d67e..a2d3b04ef7 100644 --- a/newton/_src/solvers/kamino/tests/test_solvers_metrics.py +++ b/newton/_src/solvers/kamino/tests/test_solvers_metrics.py @@ -229,6 +229,47 @@ def tearDown(self): if self.verbose: msg.reset_log_level() + def _evaluate_contact_residuals(self, contacts: list[tuple[int, int, float]]) -> tuple[np.ndarray, np.ndarray]: + """Evaluate contact residuals from explicitly supplied signed distances.""" + + def build_two_boxes_on_planes(): + builder = build_box_on_plane() + return build_box_on_plane(builder=builder) + + test = TestSetup( + builder_fn=build_two_boxes_on_planes, + max_world_contacts=4, + gravity=False, + perturb=False, + device=self.default_device, + ) + + num_contacts = len(contacts) + wid = test.contacts.wid.numpy() + cid = test.contacts.cid.numpy() + gapfunc = test.contacts.gapfunc.numpy() + wid[:num_contacts] = [contact[0] for contact in contacts] + cid[:num_contacts] = [contact[1] for contact in contacts] + gapfunc[:num_contacts] = [(0.0, 0.0, 1.0, contact[2]) for contact in contacts] + + test.contacts.model_active_contacts.assign(np.array([num_contacts], dtype=np.int32)) + test.contacts.world_active_contacts.assign( + np.bincount([contact[0] for contact in contacts], minlength=2).astype(np.int32) + ) + test.contacts.wid.assign(wid) + test.contacts.cid.assign(cid) + test.contacts.gapfunc.assign(gapfunc) + + metrics = SolutionMetrics(model=test.model) + metrics.reset() + metrics._evaluate_constraint_violations_perf( + model=test.model, + data=test.data, + contacts=test.contacts, + ) + wp.synchronize() + return metrics.data.r_cts_contacts.numpy(), metrics.data.r_cts_contacts_argmax.numpy() + def test_00_make_default(self): """ Test creating a SolutionMetrics instance with default initialization. @@ -355,12 +396,15 @@ def test_02_evaluate_trivial_solution(self): msg.info("metrics.r_ncp_compl: %s", metrics.data.r_ncp_compl) msg.info("metrics.r_vi_natmap: %s\n", metrics.data.r_vi_natmap) - # Extract the maximum contact penetration to use for validation + # Extract the maximum unilateral penetration depth, max(0, -d). nc = test.contacts.model_active_contacts.numpy()[0] - max_contact_penetration = 0.0 - for cid in range(nc): - pen = test.contacts.gapfunc.numpy()[cid][3] - max_contact_penetration = max(max_contact_penetration, pen) + signed_distances = test.contacts.gapfunc.numpy()[:nc, 3] + contact_residuals = np.maximum(0.0, -signed_distances) + max_contact_penetration = np.max(contact_residuals, initial=0.0) + max_contact_argmax = -1 + for cid, residual in enumerate(contact_residuals): + if residual > 0.0 and residual >= max_contact_penetration: + max_contact_argmax = cid # Check that all metrics are zero np.testing.assert_allclose(metrics.data.r_eom.numpy()[0], 0.0) @@ -393,9 +437,7 @@ def test_02_evaluate_trivial_solution(self): # NOTE: all contacts will have the same residual, # so the argmax will evaluate to the last constraint np.testing.assert_allclose(metrics.data.r_v_plus_argmax.numpy()[0], 11) - # NOTE: all contacts will have the same penetration, - # so the argmax will evaluate to the last contact - np.testing.assert_allclose(metrics.data.r_cts_contacts_argmax.numpy()[0], 3) + np.testing.assert_allclose(metrics.data.r_cts_contacts_argmax.numpy()[0], max_contact_argmax) np.testing.assert_allclose(metrics.data.r_ncp_primal_argmax.numpy()[0], 3) np.testing.assert_allclose(metrics.data.r_ncp_dual_argmax.numpy()[0], 3) np.testing.assert_allclose(metrics.data.r_ncp_compl_argmax.numpy()[0], 3) @@ -459,12 +501,10 @@ def test_03_evaluate_padmm_solution_box_on_plane(self): msg.info("metrics.r_ncp_compl: %s", metrics.data.r_ncp_compl) msg.info("metrics.r_vi_natmap: %s\n", metrics.data.r_vi_natmap) - # Extract the maximum contact penetration to use for validation + # Extract the maximum unilateral penetration depth, max(0, -d). nc = test.contacts.model_active_contacts.numpy()[0] - max_contact_penetration = 0.0 - for cid in range(nc): - pen = test.contacts.gapfunc.numpy()[cid][3] - max_contact_penetration = max(max_contact_penetration, pen) + signed_distances = test.contacts.gapfunc.numpy()[:nc, 3] + max_contact_penetration = np.max(np.maximum(0.0, -signed_distances), initial=0.0) # Check that all metrics are zero accuracy = 5 # number of decimal places for accuracy @@ -547,11 +587,9 @@ def test_04_evaluate_padmm_solution_boxes_hinged(self): msg.info("metrics.r_ncp_compl: %s", metrics.data.r_ncp_compl) msg.info("metrics.r_vi_natmap: %s\n", metrics.data.r_vi_natmap) - # Extract the maximum contact penetration to use for validation - max_contact_penetration = 0.0 - for cid in range(nc): - pen = test.contacts.gapfunc.numpy()[cid][3] - max_contact_penetration = max(max_contact_penetration, pen) + # Extract the maximum unilateral penetration depth, max(0, -d). + signed_distances = test.contacts.gapfunc.numpy()[:nc, 3] + max_contact_penetration = np.max(np.maximum(0.0, -signed_distances), initial=0.0) # Check that all metrics are zero accuracy = 5 # number of decimal places for accuracy @@ -823,6 +861,32 @@ def perturb_array(arr: wp.array[wp.float32]): metrics_dense.data.r_vi_natmap.numpy(), metrics_sparse.data.r_vi_natmap.numpy(), rtol=rtol, atol=atol ) + def test_07_contact_residual_positive_gaps(self): + residual, argmax = self._evaluate_contact_residuals( + [ + (0, 0, 0.0), + (0, 1, 0.1), + (1, 0, 0.2), + ] + ) + + np.testing.assert_allclose(residual, [0.0, 0.0]) + np.testing.assert_array_equal(argmax, [-1, -1]) + + def test_08_contact_residual_mixed_signed_distances(self): + residual, argmax = self._evaluate_contact_residuals( + [ + (0, 0, 0.5), + (0, 1, -0.1), + (1, 0, -0.2), + (1, 1, 0.3), + (1, 2, -0.4), + ] + ) + + np.testing.assert_allclose(residual, [0.1, 0.4]) + np.testing.assert_array_equal(argmax, [1, 2]) + ### # Test execution diff --git a/newton/_src/solvers/kamino/tests/test_solvers_padmm.py b/newton/_src/solvers/kamino/tests/test_solvers_padmm.py index 8f2d6fee6c..7693109a05 100644 --- a/newton/_src/solvers/kamino/tests/test_solvers_padmm.py +++ b/newton/_src/solvers/kamino/tests/test_solvers_padmm.py @@ -9,7 +9,6 @@ import warp as wp from newton._src.solvers.kamino._src.core.builder import ModelBuilderKamino -from newton._src.solvers.kamino._src.core.math import screw from newton._src.solvers.kamino._src.core.model import ModelKamino from newton._src.solvers.kamino._src.dynamics.dual import DualProblem from newton._src.solvers.kamino._src.kinematics.constraints import unpack_constraint_solutions @@ -50,9 +49,10 @@ def __init__( self.builder: ModelBuilderKamino = builder_fn(**kwargs) # Set ad-hoc configurations - self.builder.gravity[0].enabled = gravity + if not gravity: + self.builder.set_gravity(wp.vec3f(0.0)) if perturb: - u_0 = screw(wp.vec3f(+10.0, 0.0, 0.0), wp.vec3f(0.0, 0.0, 0.0)) + u_0 = wp.spatial_vectorf(10.0, 0.0, 0.0, 0.0, 0.0, 0.0) for body in self.builder.all_bodies: body.u_i_0 = u_0 @@ -181,8 +181,10 @@ def check_padmm_solution( test.assertLessEqual(r_p, solver.config[w].primal_tolerance) test.assertLessEqual(r_d, solver.config[w].dual_tolerance) test.assertLessEqual(r_c, solver.config[w].compl_tolerance) - test.assertLessEqual(error_dual_abs_l2, solver.config[w].dual_tolerance) - test.assertLessEqual(error_dual_abs_inf, solver.config[w].dual_tolerance) + # Using expanded tolerance for true dual error due to potential numerical inaccuracies + # between in-solver residual and true residual. + test.assertLessEqual(error_dual_abs_l2, solver.config[w].dual_tolerance * 4.0) + test.assertLessEqual(error_dual_abs_inf, solver.config[w].dual_tolerance * 4.0) def save_solver_info(solver: PADMMSolver, path: str | None = None, verbose: bool = False): diff --git a/newton/_src/solvers/kamino/tests/test_utils_io_usd.py b/newton/_src/solvers/kamino/tests/test_utils_io_usd.py index 778d7badce..35afa82bc4 100644 --- a/newton/_src/solvers/kamino/tests/test_utils_io_usd.py +++ b/newton/_src/solvers/kamino/tests/test_utils_io_usd.py @@ -11,9 +11,11 @@ import newton from newton import Model, ModelBuilder +from newton._src.core.types import Axis from newton._src.geometry.types import GeoType from newton._src.solvers.kamino import SolverKamino from newton._src.solvers.kamino._src.core.builder import ModelBuilderKamino +from newton._src.solvers.kamino._src.core.gravity import GravityDescriptor from newton._src.solvers.kamino._src.core.joints import JOINT_QMAX, JOINT_QMIN, JointActuationType, JointDoFType from newton._src.solvers.kamino._src.models.builders import basics from newton._src.solvers.kamino._src.utils import logger as msg @@ -47,6 +49,54 @@ def tearDown(self): if self.verbose: msg.reset_log_level() + def test_gravity_descriptor_from_usd_default_magnitude(self): + """Resolve OpenUSD's negative-infinity gravity sentinel.""" + gravity = GravityDescriptor.from_usd((0.0, 0.0, 0.0), -float("inf"), Axis.Y, 1.0) + + np.testing.assert_array_equal(gravity.vector, np.array([0.0, -9.81, 0.0], dtype=np.float32)) + + def test_gravity_descriptor_from_usd_negative_magnitude(self): + """Preserve an explicitly authored negative gravity magnitude.""" + gravity = GravityDescriptor.from_usd((0.0, 0.0, -1.0), -1.0, Axis.Y, 1.0) + + np.testing.assert_array_equal(gravity.vector, np.array([0.0, 0.0, 1.0], dtype=np.float32)) + + def test_gravity_descriptor_from_usd_explicit_values(self): + """Normalize and scale explicitly authored OpenUSD gravity.""" + gravity = GravityDescriptor.from_usd((3.0, 4.0, 0.0), 8.0, Axis.Z, 0.5) + + np.testing.assert_allclose(gravity.vector, np.array([2.4, 3.2, 0.0], dtype=np.float32)) + + @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core") + def test_import_default_physics_scene_gravity(self): + """Import the resolved default gravity of a USD physics scene.""" + from pxr import Usd, UsdGeom, UsdPhysics + + stage = Usd.Stage.CreateInMemory() + UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.y) + UsdGeom.SetStageMetersPerUnit(stage, 1.0) + UsdPhysics.Scene.Define(stage, "/PhysicsScene") + + builder = USDImporter().import_from(stage, load_static_geometry=False, load_materials=False) + + np.testing.assert_array_equal(builder.gravity[0].vector, np.array([0.0, -9.81, 0.0], dtype=np.float32)) + + @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core") + def test_import_zero_gravity_uses_stage_up_axis(self): + """Retain the stage up axis when imported gravity is zero.""" + from pxr import Usd, UsdGeom, UsdPhysics + + stage = Usd.Stage.CreateInMemory() + UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.y) + UsdGeom.SetStageMetersPerUnit(stage, 1.0) + scene = UsdPhysics.Scene.Define(stage, "/PhysicsScene") + scene.CreateGravityMagnitudeAttr(0.0) + + builder = USDImporter().import_from(stage, load_static_geometry=False, load_materials=False) + + self.assertEqual(builder.up_axes[0], Axis.Y) + np.testing.assert_array_equal(builder.gravity[0].vector, np.zeros(3, dtype=np.float32)) + ### # Joints supported natively by USD ### diff --git a/newton/_src/solvers/kamino/tests/utils/checks.py b/newton/_src/solvers/kamino/tests/utils/checks.py index 75274d65d0..9d1f8d1ab8 100644 --- a/newton/_src/solvers/kamino/tests/utils/checks.py +++ b/newton/_src/solvers/kamino/tests/utils/checks.py @@ -282,7 +282,6 @@ def assert_builders_equal( for m in range(builder1.num_materials): test.assertEqual(builder1.materials[m].wid, builder2.materials[m].wid) test.assertEqual(builder1.materials[m].mid, builder2.materials[m].mid) - test.assertEqual(builder1.materials[m].density, builder2.materials[m].density) test.assertEqual(builder1.materials[m].restitution, builder2.materials[m].restitution) test.assertEqual(builder1.materials[m].static_friction, builder2.materials[m].static_friction) test.assertEqual(builder1.materials[m].dynamic_friction, builder2.materials[m].dynamic_friction) @@ -524,7 +523,6 @@ def assert_model_materials_equal( ) -> None: assert_scalar_attributes_equal(test, materials0, materials1, ["num_materials"]) array_attributes = [ - # "density", "restitution", "static_friction", "dynamic_friction", diff --git a/newton/_src/solvers/kamino/tests/utils/make.py b/newton/_src/solvers/kamino/tests/utils/make.py index 86d223560b..9e07187730 100644 --- a/newton/_src/solvers/kamino/tests/utils/make.py +++ b/newton/_src/solvers/kamino/tests/utils/make.py @@ -14,7 +14,7 @@ from ..._src.core.bodies import update_body_inertias from ..._src.core.builder import ModelBuilderKamino from ..._src.core.data import DataKamino -from ..._src.core.math import quat_exp, screw, screw_angular, screw_linear +from ..._src.core.math import quat_exp from ..._src.core.model import ModelKamino from ..._src.core.state import StateKamino from ..._src.geometry.contacts import ContactsKamino @@ -337,8 +337,8 @@ def _set_fourbar_body_states( R_B = wp.quat_to_matrix(q_B) # Extract the linear and angular velocity of the Base body - v_B = screw_linear(u_B) - omega_B = screw_angular(u_B) + v_B = wp.spatial_top(u_B) + omega_B = wp.spatial_bottom(u_B) # Define the joint rotation offset q_x_j = Q_X_J * wp.pow(-1.0, float(jid)) # Alternate sign for each joint @@ -372,7 +372,7 @@ def _set_fourbar_body_states( # Offset the pose of the body by a fixed amount state_body_q_i[bid_F] = wp.transformation(r_F_new, q_F_new, dtype=wp.float32) - state_body_u_i[bid_F] = screw(v_F_new, omega_F_new) + state_body_u_i[bid_F] = wp.spatial_vectorf(*v_F_new, *omega_F_new) def set_fourbar_body_states(model: ModelKamino, data: DataKamino): diff --git a/newton/_src/solvers/kamino/tests/utils/sampling.py b/newton/_src/solvers/kamino/tests/utils/sampling.py index d28c43e83f..14bca111eb 100644 --- a/newton/_src/solvers/kamino/tests/utils/sampling.py +++ b/newton/_src/solvers/kamino/tests/utils/sampling.py @@ -66,7 +66,7 @@ def sample_base_state( max_lin_vel: float = 0.5, max_ang_vel: float = np.radians(90.0), unit_quaternions: bool = True, -) -> np.ndarray: +) -> tuple[np.ndarray, np.ndarray]: """ Helper sampling random base_q, base_u given the number of worlds. diff --git a/newton/_src/solvers/mujoco/kernels.py b/newton/_src/solvers/mujoco/kernels.py index 33842ead32..9ffae07ebf 100644 --- a/newton/_src/solvers/mujoco/kernels.py +++ b/newton/_src/solvers/mujoco/kernels.py @@ -418,6 +418,9 @@ def convert_newton_contacts_to_mjwarp_kernel( rigid_contact_damping: wp.array[wp.float32], rigid_contact_friction: wp.array[wp.float32], shape_margin: wp.array[float], + shape_material_kf: wp.array[float], + opt_impratio_invsqrt: wp.array[float], + use_kf_mapping: bool, bodies_per_world: int, newton_shape_to_mjc_geom: wp.array[wp.int32], # Mujoco warp contacts @@ -626,6 +629,26 @@ def convert_newton_contacts_to_mjwarp_kernel( friction[4], ) + # Match Newton's force-space friction slope using MuJoCo's inverse-weight + # approximation; positive solref lets refsafe limit overly stiff damping. + if shape_material_kf and use_kf_mapping: + kf1 = shape_material_kf[shape_a] + kf2 = shape_material_kf[shape_b] + kf = mix * kf1 + (1.0 - mix) * kf2 + if kf > 0.0: + invw = body_invweight0[worldid, mj_body_a][0] + body_invweight0[worldid, mj_body_b][0] + ir = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]] + imp = solimp[1] + denom = kf * invw * ((1.0 - imp) * ir * ir + imp) + if denom > 0.0 and wp.isfinite(denom): + timeconst = 2.0 / denom + if wp.isfinite(timeconst): + solreffriction = wp.vec2(timeconst, 1.0) + elif kf == 0.0: + # A zero gain means no friction force in Newton, so omit all + # sliding, torsional, and rolling constraint rows. + condim = 1 + cid = wp.atomic_add(nacon_out, 0, 1) if cid >= naconmax: tid_to_cid[tid] = -1 @@ -2490,6 +2513,54 @@ def update_geom_properties_kernel( geom_quat[world, geom_idx] = quat_xyzw_to_wxyz(tf.q) +@wp.kernel +def update_site_properties_kernel( + shape_transform: wp.array[wp.transform], + shape_scale: wp.array[wp.vec3], + site_shape_index: wp.array[wp.int32], + site_is_global: wp.array[bool], + shapes_per_world: int, + first_env_shape_base: int, + site_pos: wp.array2d[wp.vec3], + site_quat: wp.array2d[wp.quat], + site_size: wp.array[wp.vec3], +): + """Update MuJoCo site poses and sizes from Newton shape properties.""" + world, site = wp.tid() + template_or_global_shape = site_shape_index[site] + if template_or_global_shape < 0: + return + + shape = template_or_global_shape + if not site_is_global[site]: + shape = first_env_shape_base + template_or_global_shape + world * shapes_per_world + + tf = shape_transform[shape] + site_pos[world, site] = tf.p + site_quat[world, site] = quat_xyzw_to_wxyz(tf.q) + + # site_size has no world dimension in mujoco_warp, so world 0 is the + # source of truth; per-world site sizes are not representable. + if world == 0: + scale = shape_scale[shape] + # Mirror export: fill zero components with the first positive one. + nonzero = 0.0 + if scale[0] > 0.0: + nonzero = scale[0] + elif scale[1] > 0.0: + nonzero = scale[1] + elif scale[2] > 0.0: + nonzero = scale[2] + if nonzero > 0.0: + site_size[site] = wp.vec3( + wp.where(scale[0] == 0.0, nonzero, scale[0]), + wp.where(scale[1] == 0.0, nonzero, scale[1]), + wp.where(scale[2] == 0.0, nonzero, scale[2]), + ) + else: + site_size[site] = wp.vec3(0.01, 0.01, 0.01) + + @wp.kernel def sync_worldbody_geom_xposes_kernel( geom_bodyid: wp.array[int], @@ -2508,6 +2579,25 @@ def sync_worldbody_geom_xposes_kernel( geom_xmat[world, geom] = wp.quat_to_matrix(geom_q) +@wp.kernel +def sync_site_xposes_kernel( + site_bodyid: wp.array[int], + site_pos: wp.array2d[wp.vec3], + site_quat: wp.array2d[wp.quat], + body_xpos: wp.array2d[wp.vec3], + body_xquat: wp.array2d[wp.quat], + site_xpos: wp.array2d[wp.vec3], + site_xmat: wp.array2d[wp.mat33], +): + """Refresh derived site poses after per-world model updates.""" + world, site = wp.tid() + body = site_bodyid[site] + body_q = quat_wxyz_to_xyzw(body_xquat[world, body]) + site_q = quat_wxyz_to_xyzw(site_quat[world, site]) + site_xpos[world, site] = body_xpos[world, body] + wp.quat_rotate(body_q, site_pos[world, site]) + site_xmat[world, site] = wp.quat_to_matrix(body_q * site_q) + + @wp.kernel def update_jnt_solref_from_invweight0_kernel( mjc_jnt_to_newton_dof: wp.array2d[wp.int32], diff --git a/newton/_src/solvers/mujoco/solver_mujoco.py b/newton/_src/solvers/mujoco/solver_mujoco.py index c5617c9ae5..c421a4a196 100644 --- a/newton/_src/solvers/mujoco/solver_mujoco.py +++ b/newton/_src/solvers/mujoco/solver_mujoco.py @@ -78,6 +78,7 @@ reset_joint_state_kernel, reset_world_buffers_kernel, sync_qpos0_kernel, + sync_site_xposes_kernel, sync_worldbody_geom_xposes_kernel, update_axis_properties_kernel, update_body_inertia_kernel, @@ -100,6 +101,7 @@ update_model_properties_kernel, update_pair_properties_kernel, update_shape_mappings_kernel, + update_site_properties_kernel, update_solver_options_kernel, update_tendon_properties_kernel, ) @@ -3315,7 +3317,6 @@ def __init__( disable_contacts: bool = False, update_data_interval: int = 1, save_to_mjcf: str | None = None, - ls_parallel: bool | None = None, # Deprecated: ignored since mujoco_warp 3.9.1 use_mujoco_contacts: bool = True, include_sites: bool = True, skip_visual_only_geoms: bool = True, @@ -3355,7 +3356,6 @@ def __init__( disable_contacts: If True, disable contact computation in MuJoCo. update_data_interval: Frequency (in simulation steps) at which to update the MuJoCo Data object from the Newton state. If 0, Data is never updated after initialization. save_to_mjcf: Optional path to save the generated MJCF model file. - ls_parallel: Deprecated and ignored. Parallel line search was removed from ``mujoco_warp`` in 3.9.1; passing this option emits a ``DeprecationWarning`` and has no effect. use_mujoco_contacts: If True, use the MuJoCo contact solver. If False, use the Newton contact solver (newton contacts must be passed in through the step function in that case). include_sites: If ``True`` (default), Newton shapes marked with ``ShapeFlags.SITE`` are exported as MuJoCo sites. Sites are non-colliding reference points used for sensor attachment, debugging, or as frames of reference. If ``False``, sites are skipped during export. Defaults to ``True``. skip_visual_only_geoms: If ``True`` (default), geometries used only for visualization (i.e. not involved in collision) are excluded from the exported MuJoCo spec. This avoids mismatches with models that use explicit ```` definitions for collision geometry. @@ -3363,14 +3363,6 @@ def __init__( :class:`warp.DeterministicMode`, or ``None`` to inherit ``wp.config.deterministic``. """ - if ls_parallel is not None: - warnings.warn( - "ls_parallel is deprecated and no longer has any effect: parallel " - "line search was removed from mujoco_warp in 3.9.1.", - DeprecationWarning, - stacklevel=2, - ) - super().__init__(model) # Import and cache MuJoCo modules (only happens once per class) @@ -3395,6 +3387,9 @@ def __init__( """Mapping from MuJoCo [world, body] to Newton body index. Shape [nworld, nbody], dtype int32.""" self.mjc_geom_to_newton_shape: wp.array2d[wp.int32] | None = None """Mapping from MuJoCo [world, geom] to Newton shape index. Shape [nworld, ngeom], dtype int32.""" + # Template-relative for per-world sites and absolute for global sites. + self._mjc_site_shape_index: wp.array[wp.int32] | None = None + self._mjc_site_is_global: wp.array[bool] | None = None self.mjc_jnt_to_newton_jnt: wp.array2d[wp.int32] | None = None """Mapping from MuJoCo [world, joint] to Newton joint index. Shape [nworld, njnt], dtype int32.""" self.mjc_jnt_to_newton_dof: wp.array2d[wp.int32] | None = None @@ -4019,6 +4014,9 @@ def _convert_contacts_to_mjwarp(self, model: Model, state_in: State, contacts: C contacts.rigid_contact_damping, contacts.rigid_contact_friction, model.shape_margin, + model.shape_material_kf, + self.mjw_model.opt.impratio_invsqrt, + self.mjw_model.opt.cone == self._mujoco.mjtCone.mjCONE_ELLIPTIC, bodies_per_world, self.newton_shape_to_mjc_geom, # Mujoco warp contacts @@ -4138,6 +4136,7 @@ def _notify_model_changed(self, flags: ModelFlags | int) -> None: need_length_range = True if flags & ModelFlags.SHAPE_PROPERTIES: self._update_geom_properties() + self._update_site_properties() self._update_pair_properties() self._invalidate_contact_fast_path() if flags & ModelFlags.MODEL_PROPERTIES: @@ -4236,6 +4235,7 @@ def _notify_model_changed(self, flags: ModelFlags | int) -> None: if flags & ModelFlags.SHAPE_PROPERTIES: self._sync_worldbody_geom_xposes() + self._sync_site_xposes() def _sync_equality_properties_to_mujoco_cpu(self) -> None: """Mirror equality properties from MJWarp buffers to MuJoCo-C CPU buffers.""" @@ -4272,6 +4272,27 @@ def _sync_worldbody_geom_xposes(self) -> None: device=self.model.device, ) + def _sync_site_xposes(self) -> None: + """Refresh derived site poses after per-world model updates.""" + if self.mj_model.nsite == 0: + return + wp.launch( + sync_site_xposes_kernel, + dim=(self.mjw_data.nworld, self.mj_model.nsite), + inputs=[ + self.mjw_model.site_bodyid, + self.mjw_model.site_pos, + self.mjw_model.site_quat, + self.mjw_data.xpos, + self.mjw_data.xquat, + ], + outputs=[ + self.mjw_data.site_xpos, + self.mjw_data.site_xmat, + ], + device=self.model.device, + ) + def _create_inverse_shape_mapping(self): """ Create the inverse shape mapping (Newton shape -> MuJoCo [world, geom]). @@ -6585,6 +6606,12 @@ def add_body_equality(i: int): shape_to_geom_idx[shape] = geom_idx geom_to_shape_idx[geom_idx] = shape + site_to_shape_idx = {} + for shape, site_name in site_mapping.items(): + site_idx = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SITE, site_name) + if site_idx >= 0: + site_to_shape_idx[site_idx] = shape + with wp.ScopedDevice(model.device): # create the MuJoCo Warp model self.mjw_model = mujoco_warp.put_model(self.mj_model) @@ -6645,6 +6672,18 @@ def add_body_equality(i: int): device=model.device, ) + site_to_shape_idx_np = np.full((self.mj_model.nsite,), -1, dtype=np.int32) + site_is_global_np = np.zeros((self.mj_model.nsite,), dtype=bool) + for site_idx, abs_shape_idx in site_to_shape_idx.items(): + if shape_world[abs_shape_idx] < 0: + site_to_shape_idx_np[site_idx] = abs_shape_idx + site_is_global_np[site_idx] = True + else: + site_to_shape_idx_np[site_idx] = abs_shape_idx - first_env_shape_base + + self._mjc_site_shape_index = wp.array(site_to_shape_idx_np, dtype=wp.int32) + self._mjc_site_is_global = wp.array(site_is_global_np, dtype=bool) + # Create mjc_body_to_newton: MuJoCo[world, body] -> Newton body # body_mapping is {newton_body_id: mjc_body_id}, we need to invert it # and expand to 2D for all worlds @@ -6968,8 +7007,8 @@ def _expand_model_fields(self, mj_model: MjWarpModel, nworld: int): "geom_margin", "geom_gap", # "geom_rgba", - # "site_pos", - # "site_quat", + "site_pos", + "site_quat", # "cam_pos", # "cam_quat", # "cam_poscom0", @@ -7874,6 +7913,34 @@ def _update_geom_properties(self): device=self.model.device, ) + def _update_site_properties(self) -> None: + """Update MuJoCo site poses and sizes from Newton shape properties. + + ``site_size`` is unbatched in mujoco_warp, so sizes are synced from + the first world; per-world scale differences are not supported. + """ + if self.mj_model.nsite == 0: + return + + wp.launch( + update_site_properties_kernel, + dim=(self.mjw_data.nworld, self.mj_model.nsite), + inputs=[ + self.model.shape_transform, + self.model.shape_scale, + self._mjc_site_shape_index, + self._mjc_site_is_global, + self._shapes_per_world, + self._first_env_shape_base, + ], + outputs=[ + self.mjw_model.site_pos, + self.mjw_model.site_quat, + self.mjw_model.site_size, + ], + device=self.model.device, + ) + def _update_solref_from_invweight0(self): """Scale joint-limit ``jnt_solref`` using ``dof_invweight0`` and ``jnt_solimp``. diff --git a/newton/_src/solvers/semi_implicit/kernels_body.py b/newton/_src/solvers/semi_implicit/kernels_body.py index 8a151d042c..4845240e95 100644 --- a/newton/_src/solvers/semi_implicit/kernels_body.py +++ b/newton/_src/solvers/semi_implicit/kernels_body.py @@ -203,10 +203,7 @@ def eval_body_joints( axis_p = wp.transform_vector(X_wp, axis) axis_c = wp.transform_vector(X_wc, axis) - # swing twist decomposition - twist = wp.quat_twist(axis, r_err) - - q = wp.acos(twist[3]) * 2.0 * wp.sign(wp.dot(axis, wp.vec3(twist[0], twist[1], twist[2]))) + q = wp.quat_twist_angle_signed(axis, r_err) qd = wp.dot(w_err, axis_p) t_total = axis_p * ( @@ -346,10 +343,7 @@ def eval_body_joints( axis_p = wp.transform_vector(X_wp, axis) axis_c = wp.transform_vector(X_wc, axis) - # swing twist decomposition - twist = wp.quat_twist(axis, r_err) - - q = wp.acos(twist[3]) * 2.0 * wp.sign(wp.dot(axis, wp.vec3(twist[0], twist[1], twist[2]))) + q = wp.quat_twist_angle_signed(axis, r_err) qd = wp.dot(w_err, axis_p) t_total = axis_p * ( diff --git a/newton/_src/solvers/vbd/rigid_vbd_kernels.py b/newton/_src/solvers/vbd/rigid_vbd_kernels.py index 2c0a111cf2..30914e3fd6 100644 --- a/newton/_src/solvers/vbd/rigid_vbd_kernels.py +++ b/newton/_src/solvers/vbd/rigid_vbd_kernels.py @@ -51,11 +51,28 @@ _NUM_CONTACT_THREADS_PER_BODY = wp.constant(4) """Threads per body for contact accumulation using strided iteration""" -_STICK_FLAG_ANCHOR = wp.constant(1) -"""contact_stick_flag value: frozen anchor (sticking kinematic/static contacts)""" +# DER bend-twist strain measure tolerances (curvature binormal + Bishop transport). +_CABLE_KB_FOLD_EPS = wp.constant(1.0e-12) +"""Degenerate-fold scale and denominator floor for the DER curvature binormal. -_STICK_FLAG_DEADZONE = wp.constant(2) -"""contact_stick_flag value: anti-creep deadzone (sticking dynamic-dynamic contacts)""" +When both 1 + dot(t0, t1) and |cross(t0, t1)|^2 fall below this value, the +tangents are treated as a true fold with a chosen perpendicular direction. This +is only a divide-by-zero guard; magnitude is bounded by _CABLE_KB_CURVATURE_CAP.""" + +_CABLE_KB_CURVATURE_CAP = wp.constant(20.0) +"""Numerical cap for near-fold DER curvature-binormal magnitude. + +For |kb| = 2*tan(theta/2), this starts near theta ~= 168.6 deg; it is a +conditioning guard, not a material parameter.""" + +_CABLE_TRANSPORT_DENOM_EPS = wp.constant(1.0e-8) +"""Near-anti-parallel threshold for switching closed-form transport to Bishop. + +This is larger than _CABLE_KB_FOLD_EPS because transport has no curvature cap; +the closed-form expression must be left before it becomes ill-conditioned.""" + +_CABLE_TWIST_ATAN2_DENOM_EPS = wp.constant(1.0e-12) +"""Floor on sin^2 + cos^2 in the transported-twist atan2 derivative.""" # --------------------------------- # Helper classes and device functions @@ -65,12 +82,7 @@ @wp.struct class RigidContactHistory: lambda_: wp.array[wp.vec3] - stick_flag: wp.array[wp.int32] penalty_k: wp.array[float] - point0: wp.array[wp.vec3] - point1: wp.array[wp.vec3] - offset0: wp.array[wp.vec3] - offset1: wp.array[wp.vec3] normal: wp.array[wp.vec3] @@ -218,7 +230,7 @@ def ldlt6_solve(h_ll: wp.mat33, h_aa: wp.mat33, h_al: wp.mat33, rhs_lin: wp.vec3 @wp.func def compute_kappa(q_wp: wp.quat, q_wc: wp.quat, q_wp_rest: wp.quat, q_wc_rest: wp.quat) -> wp.vec3: - """Compute cable bending curvature vector kappa in the parent frame. + """Compute rest-relative angular rotation vector kappa in the parent frame. Kappa is the rotation vector (theta*axis) from the rest-aligned relative rotation. @@ -229,7 +241,7 @@ def compute_kappa(q_wp: wp.quat, q_wc: wp.quat, q_wp_rest: wp.quat, q_wc_rest: w q_wc_rest: Child rest orientation (world). Returns: - wp.vec3: Curvature vector kappa in parent frame (rotation vector form). + wp.vec3: Rotation vector kappa in parent frame. """ # Build R_align = R_rel * R_rel_rest^T using quaternions q_rel = wp.quat_inverse(q_wp) * q_wc @@ -245,6 +257,422 @@ def compute_kappa(q_wp: wp.quat, q_wc: wp.quat, q_wp_rest: wp.quat, q_wc_rest: w return axis * angle +@wp.func +def _quat_rotate_local_z(q: wp.quat) -> wp.vec3: + """Rotate local +Z by a unit quaternion; the third rotation-matrix column.""" + x = q[0] + y = q[1] + z = q[2] + w = q[3] + return wp.vec3(2.0 * (x * z + y * w), 2.0 * (y * z - x * w), 1.0 - 2.0 * (x * x + y * y)) + + +@wp.func +def _quat_rotate_local_x(q: wp.quat) -> wp.vec3: + """Rotate local +X by a unit quaternion; the first rotation-matrix column.""" + x = q[0] + y = q[1] + z = q[2] + w = q[3] + return wp.vec3(1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y + z * w), 2.0 * (x * z - y * w)) + + +@wp.func +def _normalize_with_fallback(v: wp.vec3, fallback: wp.vec3) -> wp.vec3: + """Normalize v, then fallback, with a final fixed axis if both are tiny.""" + v_len = wp.length(v) + if v_len > _SMALL_LENGTH_EPS: + return v / v_len + fb_len = wp.length(fallback) + if fb_len > _SMALL_LENGTH_EPS: + return fallback / fb_len + return wp.vec3(1.0, 0.0, 0.0) + + +@wp.func +def _project_perp(v: wp.vec3, axis: wp.vec3) -> wp.vec3: + """Project v onto the plane orthogonal to an already-unit axis.""" + return v - axis * wp.dot(v, axis) + + +@wp.func +def _perpendicular_axis_with_fallback(axis: wp.vec3, preferred: wp.vec3) -> wp.vec3: + """Return a unit vector perpendicular to axis, preferring preferred when possible.""" + axis = _normalize_with_fallback(axis, wp.vec3(0.0, 0.0, 1.0)) + perp = _project_perp(preferred, axis) + perp_len = wp.length(perp) + if perp_len > _SMALL_LENGTH_EPS: + return perp / perp_len + + alt = wp.vec3(1.0, 0.0, 0.0) + if wp.abs(axis[0]) > 0.9: + alt = wp.vec3(0.0, 1.0, 0.0) + return _normalize_with_fallback(_project_perp(alt, axis), wp.vec3(0.0, 0.0, 1.0)) + + +@wp.func +def _bishop_transport_quat(t0: wp.vec3, t1: wp.vec3, fallback_axis: wp.vec3) -> wp.quat: + """Minimal no-twist rotation that transports unit tangent t0 to unit tangent t1.""" + c = wp.clamp(wp.dot(t0, t1), -1.0, 1.0) + v = wp.cross(t0, t1) + s = wp.length(v) + if s > _SMALL_ANGLE_EPS: + return wp.quat_from_axis_angle(v / s, wp.atan2(s, c)) + if c > 0.0: + return wp.quat_identity() + + axis = _perpendicular_axis_with_fallback(t0, fallback_axis) + return wp.quat_from_axis_angle(axis, wp.pi) + + +@wp.func +def _finite_curvature_binormal(t0: wp.vec3, t1: wp.vec3, fallback_axis: wp.vec3) -> wp.vec3: + """Capped DER finite curvature binormal with a true-fold fallback.""" + tangent_dot = wp.clamp(wp.dot(t0, t1), -1.0, 1.0) + tangent_cross = wp.cross(t0, t1) + denom = 1.0 + tangent_dot + cross_sq = wp.dot(tangent_cross, tangent_cross) + + # Use exact DER while cross(t0, t1) gives a direction. At an exact fold the + # cross product vanishes, so choose a stable perpendicular cap direction. + if denom <= _CABLE_KB_FOLD_EPS and cross_sq <= _CABLE_KB_FOLD_EPS: + axis = _perpendicular_axis_with_fallback(t0, fallback_axis) + return _CABLE_KB_CURVATURE_CAP * axis + + kb = (2.0 / wp.max(_CABLE_KB_FOLD_EPS, denom)) * tangent_cross + + kb_len = wp.length(kb) + if kb_len > _CABLE_KB_CURVATURE_CAP: + kb = (_CABLE_KB_CURVATURE_CAP / kb_len) * kb + return kb + + +@wp.func +def _finite_curvature_binormal_derivative( + t0: wp.vec3, + t1: wp.vec3, + dt0: wp.vec3, + dt1: wp.vec3, +) -> wp.vec3: + """Directional derivative of kb(t0, t1) for tangent perturbations dt0, dt1. + + Pure tangent-space math: it does not know where the perturbations come from. + The caller forms each one from a world rotation as ``dti = omega_world x ti`` + to build rigid-body angular Jacobian columns. + """ + tangent_dot = wp.clamp(wp.dot(t0, t1), -1.0, 1.0) + tangent_cross = wp.cross(t0, t1) + denom = 1.0 + tangent_dot + cross_sq = wp.dot(tangent_cross, tangent_cross) + ddenom = wp.dot(dt0, t1) + wp.dot(t0, dt1) + dcross = wp.cross(dt0, t1) + wp.cross(t0, dt1) + + if denom <= _CABLE_KB_FOLD_EPS and cross_sq <= _CABLE_KB_FOLD_EPS: + # The exact-fold direction is chosen, not geometric, so it has no + # meaningful derivative; return zero to keep the model bounded. + return wp.vec3(0.0) + + denom_safe = wp.max(_CABLE_KB_FOLD_EPS, denom) + inv_denom = 1.0 / denom_safe + kb_raw = (2.0 * inv_denom) * tangent_cross + dkb_raw = (2.0 * inv_denom) * dcross + if denom > _CABLE_KB_FOLD_EPS: + dkb_raw = dkb_raw - (2.0 * ddenom * inv_denom * inv_denom) * tangent_cross + + kb_len = wp.length(kb_raw) + if kb_len > _CABLE_KB_CURVATURE_CAP: + inv_len = 1.0 / kb_len + scale = _CABLE_KB_CURVATURE_CAP * inv_len + dkb_raw = scale * (dkb_raw - kb_raw * (wp.dot(kb_raw, dkb_raw) * inv_len * inv_len)) + + return dkb_raw + + +@wp.func +def _transport_material_axis(t0: wp.vec3, t1: wp.vec3, m0: wp.vec3, fallback_axis: wp.vec3) -> wp.vec3: + """Parallel-transport a material normal from tangent t0 to tangent t1.""" + c = wp.clamp(wp.dot(t0, t1), -1.0, 1.0) + denom = 1.0 + c + # Closed-form minimal rotation t0 -> t1; assumes m0 perpendicular to t0. + if denom > _CABLE_TRANSPORT_DENOM_EPS: + w = t0 + t1 + transported = m0 - (wp.dot(m0, t1) / denom) * w + return _normalize_with_fallback(transported, fallback_axis) + + # Anti-parallel tangents (~180 deg): use a Bishop transport quaternion. + B = _bishop_transport_quat(t0, t1, fallback_axis) + return _normalize_with_fallback(wp.quat_rotate(B, m0), fallback_axis) + + +@wp.func +def _transported_twist_angle_from_material_axes( + t0: wp.vec3, + t1: wp.vec3, + m0: wp.vec3, + m1: wp.vec3, + fallback_axis: wp.vec3, +) -> float: + """Signed twist from transported parent material normal to child material normal.""" + m0_transport = _transport_material_axis(t0, t1, m0, fallback_axis) + sin_theta = wp.dot(t1, wp.cross(m0_transport, m1)) + cos_theta = wp.dot(m0_transport, m1) + return wp.atan2(sin_theta, cos_theta) + + +@wp.struct +class CableBendTwistMeasure: + """Live bend/twist geometry for one cable joint. + + Measured once per force/Hessian evaluation and reused by the residual and + analytic Jacobian columns. + """ + + t0: wp.vec3 + t1: wp.vec3 + m0: wp.vec3 + m1: wp.vec3 + kb_world: wp.vec3 + twist: float + + +@wp.func +def _measure_cable_bend_twist_z(q_wp: wp.quat, q_wc: wp.quat) -> CableBendTwistMeasure: + """Measure bend/twist for SolverVBD cables, whose material tangent is local +Z. + + The fixed cable material basis is local ``+X, +Y, +Z``. + SolverVBD keeps body rotations normalized, so rotated basis axes are already + orthonormal. + """ + t0 = _quat_rotate_local_z(q_wp) + t1 = _quat_rotate_local_z(q_wc) + m0 = _quat_rotate_local_x(q_wp) + m1 = _quat_rotate_local_x(q_wc) + + # DER-style split: bend comes from the finite curvature binormal of the two + # tangents; twist is material spin after no-twist/Bishop transport. + measure = CableBendTwistMeasure() + measure.t0 = t0 + measure.t1 = t1 + measure.m0 = m0 + measure.m1 = m1 + measure.twist = _transported_twist_angle_from_material_axes(t0, t1, m0, m1, m0) + measure.kb_world = _finite_curvature_binormal(t0, t1, m0) + return measure + + +@wp.func +def _transported_twist_angle_derivative_from_measure( + measure: CableBendTwistMeasure, + omega_world: wp.vec3, + is_parent: bool, +) -> float: + """Directional derivative of transported twist for one endpoint rotation. + + Linear in the infinitesimal world-space rotation ``omega_world``; frames + are read from ``measure``. + """ + t0 = measure.t0 + t1 = measure.t1 + m0 = measure.m0 + m1 = measure.m1 + + dt0 = wp.vec3(0.0) + dt1 = wp.vec3(0.0) + dm0 = wp.vec3(0.0) + dm1 = wp.vec3(0.0) + if is_parent: + dt0 = wp.cross(omega_world, t0) + dm0 = wp.cross(omega_world, m0) + else: + dt1 = wp.cross(omega_world, t1) + dm1 = wp.cross(omega_world, m1) + + c = wp.clamp(wp.dot(t0, t1), -1.0, 1.0) + denom = 1.0 + c + if denom <= _CABLE_TRANSPORT_DENOM_EPS: + # At a 180-degree kink the transport derivative is singular. Use bounded + # tangent spin while the residual fallback supplies the finite angle. + if is_parent: + return -wp.dot(omega_world, t0) + return wp.dot(omega_world, t1) + + w = t0 + t1 + n = wp.dot(m0, t1) + scale = n / denom + m0_transport_raw = m0 - scale * w + m0_transport = _normalize_with_fallback(m0_transport_raw, m0) + + ddenom = wp.dot(dt0, t1) + wp.dot(t0, dt1) + dn = wp.dot(dm0, t1) + wp.dot(m0, dt1) + dscale = (dn * denom - n * ddenom) / (denom * denom) + dm0_transport_raw = dm0 - dscale * w - scale * (dt0 + dt1) + + raw_len = wp.length(m0_transport_raw) + dm0_transport = wp.vec3(0.0) + if raw_len > _SMALL_LENGTH_EPS: + inv_len = 1.0 / raw_len + dm0_transport = inv_len * (dm0_transport_raw - m0_transport * wp.dot(m0_transport, dm0_transport_raw)) + + sin_theta = wp.dot(t1, wp.cross(m0_transport, m1)) + cos_theta = wp.dot(m0_transport, m1) + dsin = wp.dot(dt1, wp.cross(m0_transport, m1)) + wp.dot( + t1, wp.cross(dm0_transport, m1) + wp.cross(m0_transport, dm1) + ) + dcos = wp.dot(dm0_transport, m1) + wp.dot(m0_transport, dm1) + denom_angle = wp.max(_CABLE_TWIST_ATAN2_DENOM_EPS, sin_theta * sin_theta + cos_theta * cos_theta) + return (cos_theta * dsin - sin_theta * dcos) / denom_angle + + +@wp.func +def _cable_bend_twist_directional_derivatives_from_measure( + q_wp: wp.quat, + measure: CableBendTwistMeasure, + omega_world: wp.vec3, + is_parent: bool, +) -> tuple[wp.vec3, float]: + """Bend/twist derivatives for one endpoint angular perturbation. + + Returns ``(d_bend_local, d_twist)``: the un-projected parent-frame + curvature-binormal derivative and transported-twist derivative. For a parent + perturbation the bend term also differentiates the parent-frame coordinates + (the ``-omega_world x kb`` term). + """ + t0 = measure.t0 + t1 = measure.t1 + + dt0 = wp.vec3(0.0) + dt1 = wp.vec3(0.0) + if is_parent: + dt0 = wp.cross(omega_world, t0) + else: + dt1 = wp.cross(omega_world, t1) + + dkb_world = _finite_curvature_binormal_derivative(t0, t1, dt0, dt1) + + # A parent rotation also turns the parent frame the binormal is expressed in, + # so the parent-frame coordinates pick up the extra -omega x kb term before + # the single map into the parent frame. + dkb_parent_frame = dkb_world + if is_parent: + dkb_parent_frame = dkb_world - wp.cross(omega_world, measure.kb_world) + d_bend_local = wp.quat_rotate(wp.quat_inverse(q_wp), dkb_parent_frame) + + d_twist = _transported_twist_angle_derivative_from_measure(measure, omega_world, is_parent) + return d_bend_local, d_twist + + +@wp.func +def _geometric_cable_strain_directional_derivative_z_from_measure( + q_wp: wp.quat, + measure: CableBendTwistMeasure, + omega_world: wp.vec3, + is_parent: bool, +) -> wp.vec3: + """Directional derivative of [bend_x, bend_y, twist_z] for local +Z cables.""" + d_bend_local, d_twist = _cable_bend_twist_directional_derivatives_from_measure( + q_wp, measure, omega_world, is_parent + ) + return wp.vec3(d_bend_local[0], d_bend_local[1], d_twist) + + +@wp.func +def _cable_bend_twist_jacobian_z_from_measure( + q_wp: wp.quat, + measure: CableBendTwistMeasure, + is_parent: bool, +) -> wp.mat33: + """Jacobian of [bend_x, bend_y, twist_z] for fixed local +Z cables. + + The local residual is exactly ``[bend_x, bend_y, twist_z]``, so no bend + projector or twist-axis vector is needed in this hot path. + """ + e0 = wp.vec3(1.0, 0.0, 0.0) + e1 = wp.vec3(0.0, 1.0, 0.0) + e2 = wp.vec3(0.0, 0.0, 1.0) + + j0 = _geometric_cable_strain_directional_derivative_z_from_measure(q_wp, measure, e0, is_parent) + j1 = _geometric_cable_strain_directional_derivative_z_from_measure(q_wp, measure, e1, is_parent) + j2 = _geometric_cable_strain_directional_derivative_z_from_measure(q_wp, measure, e2, is_parent) + return wp.matrix_from_cols(j0, j1, j2) + + +@wp.func +def _wrap_principal_angle(angle: float) -> float: + """Wrap a bounded angular difference into ``[-pi, pi]``. + + Args: + angle: Difference of two principal angles. The caller guarantees a + value in ``[-2*pi, 2*pi]``. + + Returns: + Equivalent principal angular difference. + """ + if angle > wp.pi: + angle -= 2.0 * wp.pi + elif angle < -wp.pi: + angle += 2.0 * wp.pi + return angle + + +@wp.func +def _assemble_geometric_cable_kappa_z( + q_wp: wp.quat, + kb_now_world: wp.vec3, + twist_now: float, + kb_rest_local: wp.vec3, + twist_rest: float, +) -> wp.vec3: + """Assemble [bend_x, bend_y, twist_z] for SolverVBD local +Z cables.""" + bend_now_local = wp.quat_rotate(wp.quat_inverse(q_wp), kb_now_world) + bend_residual_local = bend_now_local - kb_rest_local + # In the local +Z convention, parent-frame x/y are bend and z is twist. + # Wrap the twist delta of two atan2 angles into [-pi, pi] to avoid a ~2*pi + # jump when rest/current straddle the branch cut (one revolution per joint). + twist_residual = _wrap_principal_angle(twist_now - twist_rest) + return wp.vec3(bend_residual_local[0], bend_residual_local[1], twist_residual) + + +@wp.func +def _cable_bend_twist_delta(kappa: wp.vec3, kappa_prev: wp.vec3) -> wp.vec3: + """Return a temporal cable bend/twist strain increment. + + Args: + kappa: Current ``[bend_x, bend_y, twist_z]`` strain. + kappa_prev: Previous strain in the same representation. + + Returns: + Bend increments from ordinary subtraction and the shortest signed + principal-angle increment for twist. + """ + return wp.vec3( + kappa[0] - kappa_prev[0], + kappa[1] - kappa_prev[1], + _wrap_principal_angle(kappa[2] - kappa_prev[2]), + ) + + +@wp.func +def compute_geometric_cable_kappa_cached_z( + q_wp: wp.quat, + q_wc: wp.quat, + kb_rest_local: wp.vec3, + twist_rest: float, +) -> wp.vec3: + """Geometric cable strain residual for fixed local +Z cables.""" + measure = _measure_cable_bend_twist_z(q_wp, q_wc) + return _assemble_geometric_cable_kappa_z(q_wp, measure.kb_world, measure.twist, kb_rest_local, twist_rest) + + +@wp.func +def _diag_mul_mat33(d: wp.vec3, m: wp.mat33) -> wp.mat33: + """Return diag(d) * m without building a dense diagonal matrix.""" + return wp.matrix_from_rows( + d[0] * wp.vec3(m[0, 0], m[0, 1], m[0, 2]), + d[1] * wp.vec3(m[1, 0], m[1, 1], m[1, 2]), + d[2] * wp.vec3(m[2, 0], m[2, 1], m[2, 2]), + ) + + @wp.func def compute_right_jacobian_inverse(kappa: wp.vec3) -> wp.mat33: """Inverse right Jacobian Jr^{-1}(kappa) for SO(3) rotation vectors. @@ -273,7 +701,7 @@ def compute_kappa_dot( omega_p_world: wp.vec3, omega_c_world: wp.vec3, ) -> wp.vec3: - """Time derivative of curvature vector d(kappa)/dt in parent frame. + """Time derivative of the rotation-vector residual d(kappa)/dt in parent frame. Exploits J_world^T = Jr_inv * R_align^T * R_wp^T, so kappa_dot = J_world^T * (omega_c - omega_p). @@ -284,7 +712,7 @@ def compute_kappa_dot( omega_c_world: Child angular velocity (world) [rad/s]. Returns: - wp.vec3: Curvature rate kappa_dot in parent frame [rad/s]. + wp.vec3: Rotation-vector rate kappa_dot in parent frame [rad/s]. """ return wp.transpose(J_world) * (omega_c_world - omega_p_world) @@ -296,10 +724,10 @@ def compute_kappa_and_jacobian( q_wp_rest: wp.quat, q_wc_rest: wp.quat, ): - """Compute curvature vector and world-frame Jacobian from quaternion poses. + """Compute rotation-vector residual and world-frame Jacobian from quaternion poses. Returns: - (kappa, J_world) -- curvature vector and world-frame force Jacobian. + (kappa, J_world) -- rotation vector and world-frame force Jacobian. """ q_rel = wp.quat_inverse(q_wp) * q_wc q_rel_rest = wp.quat_inverse(q_wp_rest) * q_wc_rest @@ -431,8 +859,6 @@ def evaluate_angular_constraint_force_hessian( is_parent: bool, penalty_k: float, P: wp.mat33, - sigma0: wp.vec3, - C_fric: wp.vec3, lambda_ang: wp.vec3, C0_ang: wp.vec3, alpha: float, @@ -441,23 +867,22 @@ def evaluate_angular_constraint_force_hessian( ): """Projected angular constraint force/Hessian using rotation-vector error (kappa). - Unified evaluator for all joint types. Computes constraint force and Hessian - in the constrained subspace defined by the orthogonal-complement projector P. + Generic evaluator for non-cable angular constraints. Computes force and + Hessian in the constrained subspace defined by the orthogonal-complement + projector P. Angular Dahl friction is cable-only and handled separately, so + this evaluator carries no friction term. C0 stabilization: when alpha > 0 and C0_ang is nonzero, the effective kappa is kappa - alpha*C0_ang (initial violation snapshot). Special cases by projector: - - P = I: isotropic (CABLE bend, FIXED angular) + - P = I: isotropic (FIXED angular) - P = I - a*a^T: revolute (1 free angular axis) - arbitrary P: D6 (0-3 free angular axes) - Dahl friction (sigma0, C_fric) is only valid when P = I (isotropic). - Pass vec3(0) for both when P != I. - Returns: (tau_world, H_aa, kappa, J_world) -- constraint torque and Hessian in world - frame, plus the curvature vector and world-frame Jacobian for reuse by the + frame, plus the rotation vector and world-frame Jacobian for reuse by the drive/limit block. """ inv_dt = 1.0 / dt @@ -469,19 +894,9 @@ def evaluate_angular_constraint_force_hessian( # P_ang is constant for joint angular residuals, so lambda_ang should already # be in-basis. Project here too so stale or externally edited state cannot # apply force along a free angular DOF. - f_local = penalty_k * kappa_perp + sigma0 + P * lambda_ang - - H_local = penalty_k * P + wp.mat33( - C_fric[0], - 0.0, - 0.0, - 0.0, - C_fric[1], - 0.0, - 0.0, - 0.0, - C_fric[2], - ) + f_local = penalty_k * kappa_perp + P * lambda_ang + + H_local = penalty_k * P if damping > 0.0: omega_p_world = quat_velocity(q_wp, q_wp_prev, dt) @@ -504,6 +919,59 @@ def evaluate_angular_constraint_force_hessian( return tau_world, H_aa, kappa_now_vec, J_world +@wp.func +def evaluate_cable_bend_twist_force_hessian_z( + q_wp: wp.quat, + q_wc: wp.quat, + kb_rest_local: wp.vec3, + twist_rest: float, + q_wp_prev: wp.quat, + q_wc_prev: wp.quat, + is_parent: bool, + K_elastic_diag: wp.vec3, + C0_force: wp.vec3, + sigma0: wp.vec3, + H_fric_diag: wp.vec3, + lambda_projected: wp.vec3, + K_damp_diag: wp.vec3, + damping_active: bool, + dt: float, +): + """Bend/twist torque and Hessian for SolverVBD local +Z cables. + + In the fixed cable material basis, local angular operators are diagonal: + ``[bend_x, bend_y, twist_z]``. Keep them as vec3 row scales in the hot path + instead of building dense local matrices. + """ + inv_dt = 1.0 / dt + + measure = _measure_cable_bend_twist_z(q_wp, q_wc) + kappa_now_vec = _assemble_geometric_cable_kappa_z(q_wp, measure.kb_world, measure.twist, kb_rest_local, twist_rest) + + # Bend and twist decouple in the material basis: the angular energy is a sum + # of independent quadratics in [bend_x, bend_y, twist_z], so elastic stiffness + # and the friction Hessian have no off-diagonal coupling. Carry them as vec3 + # row scales; the dense angular block reappears below via J^T diag(H) J. + f_local = wp.cw_mul(K_elastic_diag, kappa_now_vec) - C0_force + sigma0 + lambda_projected + H_local_diag = K_elastic_diag + H_fric_diag + + if damping_active: + prev_measure = _measure_cable_bend_twist_z(q_wp_prev, q_wc_prev) + kappa_prev_vec = _assemble_geometric_cable_kappa_z( + q_wp_prev, prev_measure.kb_world, prev_measure.twist, kb_rest_local, twist_rest + ) + dkappa_dt = _cable_bend_twist_delta(kappa_now_vec, kappa_prev_vec) * inv_dt + f_local = f_local + wp.cw_mul(K_damp_diag, dkappa_dt) + H_local_diag = H_local_diag + inv_dt * K_damp_diag + + J_body = _cable_bend_twist_jacobian_z_from_measure(q_wp, measure, is_parent) + # Gauss-Newton self Hessian: J^T diag(H_local_diag) J. + H_aa = wp.transpose(J_body) * _diag_mul_mat33(H_local_diag, J_body) + tau_world = -(wp.transpose(J_body) * f_local) + + return tau_world, H_aa, kappa_now_vec, J_body + + @wp.func def evaluate_linear_constraint_force_hessian( X_wp: wp.transform, @@ -525,14 +993,14 @@ def evaluate_linear_constraint_force_hessian( ): """Projected linear constraint force/Hessian for anchor coincidence. - Unified evaluator for all joint types. Computes C = x_c - x_p, projects - with P, and returns force/Hessian in world frame. + Generic evaluator for non-cable linear constraints. Computes C = x_c - x_p, + projects with P, and returns force/Hessian in world frame. C0 stabilization: when alpha > 0 and C0_lin is nonzero, the effective constraint violation is C - alpha*C0 (initial violation snapshot). Special cases by projector: - - P = I: isotropic (BALL, CABLE stretch, FIXED linear, REVOLUTE linear) + - P = I: isotropic (BALL, FIXED linear, REVOLUTE linear) - P = I - a*a^T: prismatic (1 free linear axis) - arbitrary P: D6 (0-3 free linear axes) @@ -582,6 +1050,71 @@ def evaluate_linear_constraint_force_hessian( return force, torque, H_ll, H_al, H_aa +@wp.func +def evaluate_cable_stretch_shear_force_hessian( + X_wp: wp.transform, + X_wc: wp.transform, + X_wp_prev: wp.transform, + X_wc_prev: wp.transform, + parent_pose: wp.transform, + child_pose: wp.transform, + parent_com: wp.vec3, + child_com: wp.vec3, + is_parent: bool, + k_diag: wp.vec3, + C0_force_local: wp.vec3, + lambda_local: wp.vec3, + kd_diag: wp.vec3, + damping_active: bool, + dt: float, +): + """Cable stretch/shear anchor force, torque, and PSD Gauss-Newton self-Hessian. + + All inputs are parent-material: residual ``u = R_p^T (x_c - x_p) = + [shear_x, shear_y, stretch_z]``, diagonal ``k_diag = (k_shear, k_shear, + k_stretch)`` / ``kd_diag``, and local ``C0_force_local`` / ``lambda_local``. + Elastic and AL are energy gradients, damping is dissipative in ``u``, and both + bodies react at the shared child anchor, so the net internal wrench is zero. + """ + x_p = wp.transform_get_translation(X_wp) + x_c = wp.transform_get_translation(X_wc) + q_wp = wp.transform_get_rotation(X_wp) + + if is_parent: + com_w = wp.transform_point(parent_pose, parent_com) + else: + com_w = wp.transform_point(child_pose, child_com) + r = x_c - com_w + + C_vec = x_c - x_p + u = wp.quat_rotate_inv(q_wp, C_vec) + psi = wp.cw_mul(k_diag, u) - C0_force_local + lambda_local + + h_s = k_diag[0] + h_z = k_diag[2] + if damping_active: + inv_dt = 1.0 / dt + x_p_prev = wp.transform_get_translation(X_wp_prev) + x_c_prev = wp.transform_get_translation(X_wc_prev) + u_prev = wp.quat_rotate_inv(wp.transform_get_rotation(X_wp_prev), x_c_prev - x_p_prev) + psi = psi + wp.cw_mul(kd_diag, (u - u_prev) * inv_dt) + h_s = h_s + kd_diag[0] * inv_dt + h_z = h_z + kd_diag[2] * inv_dt + + f_world = wp.quat_rotate(q_wp, psi) + force = f_world if is_parent else -f_world + + t = _quat_rotate_local_z(q_wp) + K_eff = h_s * wp.identity(3, float) + (h_z - h_s) * wp.outer(t, t) + rx = wp.skew(r) + H_ll = K_eff + H_al = rx * K_eff + H_aa = wp.transpose(rx) * K_eff * rx + + torque = wp.cross(r, force) + return force, torque, H_ll, H_al, H_aa + + # --------------------------------- # Data structures # --------------------------------- @@ -1211,6 +1744,8 @@ def evaluate_joint_force_hessian( joint_X_p: wp.array[wp.transform], joint_X_c: wp.array[wp.transform], joint_axis: wp.array[wp.vec3], + joint_cable_rest_kb_local: wp.array[wp.vec3], + joint_cable_rest_twist: wp.array[float], joint_qd_start: wp.array[int], joint_target_q_start: wp.array[int], joint_constraint_start: wp.array[int], @@ -1241,12 +1776,13 @@ def evaluate_joint_force_hessian( """Compute AVBD joint force and Hessian contributions for one body. Supported joint types: CABLE, BALL, FIXED, REVOLUTE, PRISMATIC, D6. - Uses unified projector-based constraint evaluators for all joint types. + Cable uses split stretch/shear and bend/twist helpers; other joints use + projector-based linear/angular evaluators. Indexing: joint_constraint_start[j] is a solver-owned start offset into the per-constraint arrays (joint_penalty_k, joint_penalty_kd). Layout per joint type: - - CABLE: 2 scalars -> [stretch, bend] + - CABLE: 4 scalars -> [stretch, shear, bend, twist] - BALL: 1 scalar -> [linear] - FIXED: 2 scalars -> [linear, angular] - REVOLUTE: 3 scalars -> [linear, angular, ang_drive_limit] @@ -1282,65 +1818,42 @@ def evaluate_joint_force_hessian( if parent_index >= 0: parent_pose = body_q[parent_index] parent_pose_prev = body_q_prev[parent_index] - parent_pose_rest = body_q_rest[parent_index] parent_com = body_com[parent_index] else: parent_pose = wp.transform(wp.vec3(0.0), wp.quat_identity()) parent_pose_prev = parent_pose - parent_pose_rest = parent_pose parent_com = wp.vec3(0.0) child_pose = body_q[child_index] child_pose_prev = body_q_prev[child_index] - child_pose_rest = body_q_rest[child_index] child_com = body_com[child_index] X_wp = parent_pose * X_pj X_wc = child_pose * X_cj X_wp_prev = parent_pose_prev * X_pj X_wc_prev = child_pose_prev * X_cj - X_wp_rest = parent_pose_rest * X_pj - X_wc_rest = child_pose_rest * X_cj c_start = joint_constraint_start[joint_index] # Hoist quaternion extraction (shared by all angular constraints and drive/limits) q_wp = wp.transform_get_rotation(X_wp) q_wc = wp.transform_get_rotation(X_wc) - q_wp_rest = wp.transform_get_rotation(X_wp_rest) - q_wc_rest = wp.transform_get_rotation(X_wc_rest) q_wp_prev = wp.transform_get_rotation(X_wp_prev) q_wc_prev = wp.transform_get_rotation(X_wc_prev) - P_I = wp.identity(3, float) - - # Hard/soft AL gating for the linear structural slot (slot 0) - lin_lambda = wp.vec3(0.0) - lin_C0 = wp.vec3(0.0) - lin_alpha = float(0.0) - if joint_is_hard[c_start] == 1: - lin_lambda = joint_lambda_lin[joint_index] - lin_C0 = joint_C0_lin[joint_index] - lin_alpha = avbd_alpha - - # Hard/soft AL gating for the angular structural slot (slot 1) - ang_lambda = wp.vec3(0.0) - ang_C0 = wp.vec3(0.0) - ang_alpha = float(0.0) - ang_hard = 0 - if jt != JointType.BALL: - ang_hard = joint_is_hard[c_start + 1] - - if ang_hard == 1: - ang_lambda = joint_lambda_ang[joint_index] - ang_C0 = joint_C0_ang[joint_index] - ang_alpha = avbd_alpha - if jt == JointType.CABLE: - k_stretch = joint_penalty_k[c_start] - k_bend = joint_penalty_k[c_start + 1] - kd_stretch = joint_penalty_kd[c_start] - kd_bend = joint_penalty_kd[c_start + 1] + stretch_idx = c_start + shear_idx = c_start + 1 + bend_idx = c_start + 2 + twist_idx = c_start + 3 + k_stretch = joint_penalty_k[stretch_idx] + k_shear = joint_penalty_k[shear_idx] + kd_stretch = joint_penalty_kd[stretch_idx] + kd_shear = joint_penalty_kd[shear_idx] + k_bend = joint_penalty_k[bend_idx] + k_twist = joint_penalty_k[twist_idx] + kd_bend = joint_penalty_kd[bend_idx] + kd_twist = joint_penalty_kd[twist_idx] total_force = wp.vec3(0.0) total_torque = wp.vec3(0.0) @@ -1348,36 +1861,90 @@ def evaluate_joint_force_hessian( total_H_al = wp.mat33(0.0) total_H_aa = wp.mat33(0.0) - if k_bend > 0.0: - if ang_hard == 1: - sigma0 = wp.vec3(0.0) - C_fric = wp.vec3(0.0) - else: - sigma0 = joint_sigma_start[joint_index] - C_fric = joint_C_fric[joint_index] - bend_torque, bend_H_aa, _bend_kappa, _bend_J = evaluate_angular_constraint_force_hessian( + bend_stiff = k_bend > 0.0 + twist_stiff = k_twist > 0.0 + bend_active = bend_stiff or kd_bend > 0.0 + twist_active = twist_stiff or kd_twist > 0.0 + if bend_active or twist_active: + K_elastic_diag = wp.vec3(k_bend, k_bend, k_twist) + # kd_X is already 0 when its slot is inactive (X_active includes kd_X > 0), + # so use the damping coefficients directly, matching K_elastic_diag above. + K_damp_diag = wp.vec3(kd_bend, kd_bend, kd_twist) + damping_active = kd_bend > 0.0 or kd_twist > 0.0 + + sigma = wp.vec3(0.0) + H_fric_diag = wp.vec3(0.0) + lambda_projected = wp.vec3(0.0) + C0_force = wp.vec3(0.0) + dahl_sigma = joint_sigma_start[joint_index] + dahl_fric = joint_C_fric[joint_index] + bend_hard = bend_stiff and joint_is_hard[bend_idx] == 1 + twist_hard = twist_stiff and joint_is_hard[twist_idx] == 1 + lambda_ang = wp.vec3(0.0) + C0_ang = wp.vec3(0.0) + if bend_hard or twist_hard: + lambda_ang = joint_lambda_ang[joint_index] + C0_ang = joint_C0_ang[joint_index] + + if bend_hard: + lambda_projected = lambda_projected + wp.vec3(lambda_ang[0], lambda_ang[1], 0.0) + C0_force = C0_force + (k_bend * avbd_alpha) * wp.vec3(C0_ang[0], C0_ang[1], 0.0) + elif bend_stiff: + sigma = sigma + wp.vec3(dahl_sigma[0], dahl_sigma[1], 0.0) + H_fric_diag = H_fric_diag + wp.vec3(dahl_fric[0], dahl_fric[1], 0.0) + + if twist_hard: + lambda_projected = lambda_projected + wp.vec3(0.0, 0.0, lambda_ang[2]) + C0_force = C0_force + (k_twist * avbd_alpha) * wp.vec3(0.0, 0.0, C0_ang[2]) + elif twist_stiff: + sigma = sigma + wp.vec3(0.0, 0.0, dahl_sigma[2]) + H_fric_diag = H_fric_diag + wp.vec3(0.0, 0.0, dahl_fric[2]) + + cable_torque, cable_H_aa, _cable_kappa, _cable_J = evaluate_cable_bend_twist_force_hessian_z( q_wp, q_wc, - q_wp_rest, - q_wc_rest, + joint_cable_rest_kb_local[joint_index], + joint_cable_rest_twist[joint_index], q_wp_prev, q_wc_prev, is_parent_body, - k_bend, - P_I, - sigma0, - C_fric, - ang_lambda, - ang_C0, - ang_alpha, - kd_bend, + K_elastic_diag, + C0_force, + sigma, + H_fric_diag, + lambda_projected, + K_damp_diag, + damping_active, dt, ) - total_torque = total_torque + bend_torque - total_H_aa = total_H_aa + bend_H_aa - - if k_stretch > 0.0: - f_s, t_s, Hll_s, Hal_s, Haa_s = evaluate_linear_constraint_force_hessian( + total_torque = total_torque + cable_torque + total_H_aa = total_H_aa + cable_H_aa + + stretch_stiff = k_stretch > 0.0 + shear_stiff = k_shear > 0.0 + stretch_active = stretch_stiff or kd_stretch > 0.0 + shear_active = shear_stiff or kd_shear > 0.0 + if stretch_active or shear_active: + # Parent-material diagonals for local u = [shear_x, shear_y, stretch_z]. + k_diag = wp.vec3(k_shear, k_shear, k_stretch) + kd_diag = wp.vec3(kd_shear, kd_shear, kd_stretch) + damping_active = kd_stretch > 0.0 or kd_shear > 0.0 + + lambda_local = wp.vec3(0.0) + C0_force_local = wp.vec3(0.0) + stretch_hard = stretch_stiff and joint_is_hard[stretch_idx] == 1 + shear_hard = shear_stiff and joint_is_hard[shear_idx] == 1 + if stretch_hard or shear_hard: + lambda_lin = joint_lambda_lin[joint_index] + C0_lin = joint_C0_lin[joint_index] + if stretch_hard: + lambda_local = lambda_local + wp.vec3(0.0, 0.0, lambda_lin[2]) + C0_force_local = C0_force_local + (k_stretch * avbd_alpha) * wp.vec3(0.0, 0.0, C0_lin[2]) + if shear_hard: + lambda_local = lambda_local + wp.vec3(lambda_lin[0], lambda_lin[1], 0.0) + C0_force_local = C0_force_local + (k_shear * avbd_alpha) * wp.vec3(C0_lin[0], C0_lin[1], 0.0) + + f_l, t_l, Hll_l, Hal_l, Haa_l = evaluate_cable_stretch_shear_force_hessian( X_wp, X_wc, X_wp_prev, @@ -1387,23 +1954,42 @@ def evaluate_joint_force_hessian( parent_com, child_com, is_parent_body, - k_stretch, - P_I, - lin_lambda, - lin_C0, - lin_alpha, - kd_stretch, + k_diag, + C0_force_local, + lambda_local, + kd_diag, + damping_active, dt, ) - total_force = total_force + f_s - total_torque = total_torque + t_s - total_H_ll = total_H_ll + Hll_s - total_H_al = total_H_al + Hal_s - total_H_aa = total_H_aa + Haa_s + total_force = total_force + f_l + total_torque = total_torque + t_l + total_H_ll = total_H_ll + Hll_l + total_H_al = total_H_al + Hal_l + total_H_aa = total_H_aa + Haa_l return total_force, total_torque, total_H_ll, total_H_al, total_H_aa - elif jt == JointType.BALL: + P_I = wp.identity(3, float) + + # Hard/soft AL gating for the non-cable linear structural slot. + lin_lambda = wp.vec3(0.0) + lin_C0 = wp.vec3(0.0) + lin_alpha = float(0.0) + if joint_is_hard[c_start] == 1: + lin_lambda = joint_lambda_lin[joint_index] + lin_C0 = joint_C0_lin[joint_index] + lin_alpha = avbd_alpha + + # BALL has no angular structural slot; other non-cable joints do. + ang_lambda = wp.vec3(0.0) + ang_C0 = wp.vec3(0.0) + ang_alpha = float(0.0) + if jt != JointType.BALL and joint_is_hard[c_start + 1] == 1: + ang_lambda = joint_lambda_ang[joint_index] + ang_C0 = joint_C0_ang[joint_index] + ang_alpha = avbd_alpha + + if jt == JointType.BALL: k = joint_penalty_k[c_start] damping = joint_penalty_kd[c_start] if k > 0.0: @@ -1427,7 +2013,15 @@ def evaluate_joint_force_hessian( ) return _zero_force_hessian() - elif jt == JointType.FIXED: + if parent_index >= 0: + X_wp_rest = body_q_rest[parent_index] * X_pj + else: + X_wp_rest = X_pj + X_wc_rest = body_q_rest[child_index] * X_cj + q_wp_rest = wp.transform_get_rotation(X_wp_rest) + q_wc_rest = wp.transform_get_rotation(X_wc_rest) + + if jt == JointType.FIXED: k_lin = joint_penalty_k[c_start + 0] kd_lin = joint_penalty_kd[c_start + 0] if k_lin > 0.0: @@ -1469,8 +2063,6 @@ def evaluate_joint_force_hessian( is_parent_body, k_ang, P_I, - wp.vec3(0.0), - wp.vec3(0.0), ang_lambda, ang_C0, ang_alpha, @@ -1534,8 +2126,6 @@ def evaluate_joint_force_hessian( is_parent_body, k_ang, P_ang, - wp.vec3(0.0), - wp.vec3(0.0), ang_lambda, ang_C0, ang_alpha, @@ -1646,8 +2236,6 @@ def evaluate_joint_force_hessian( is_parent_body, k_ang, P_ang, - wp.vec3(0.0), - wp.vec3(0.0), ang_lambda, ang_C0, ang_alpha, @@ -1791,8 +2379,6 @@ def evaluate_joint_force_hessian( is_parent_body, k_ang, P_ang, - wp.vec3(0.0), - wp.vec3(0.0), ang_lambda, ang_C0, ang_alpha, @@ -2270,13 +2856,56 @@ def check_contact_overflow( ) +@wp.kernel +def init_cable_rest_bend_twist( + joint_type: wp.array[int], + joint_parent: wp.array[int], + joint_child: wp.array[int], + joint_X_p: wp.array[wp.transform], + joint_X_c: wp.array[wp.transform], + body_q_rest: wp.array[wp.transform], + joint_cable_rest_kb_local: wp.array[wp.vec3], + joint_cable_rest_twist: wp.array[float], +): + """Precompute cable rest angular deformation invariants.""" + j = wp.tid() + joint_cable_rest_kb_local[j] = wp.vec3(0.0) + joint_cable_rest_twist[j] = 0.0 + + if joint_type[j] != JointType.CABLE: + return + + child = joint_child[j] + if child < 0: + return + + parent = joint_parent[j] + if parent >= 0: + X_wp_rest = body_q_rest[parent] * joint_X_p[j] + else: + X_wp_rest = joint_X_p[j] + X_wc_rest = body_q_rest[child] * joint_X_c[j] + + q_wp_rest = wp.transform_get_rotation(X_wp_rest) + q_wc_rest = wp.transform_get_rotation(X_wc_rest) + + # Rest DER bend (parent-local curvature binormal) and rest twist (transported + # material spin), measured once so a pre-curved rest yields zero strain. + rest_measure = _measure_cable_bend_twist_z(q_wp_rest, q_wc_rest) + joint_cable_rest_kb_local[j] = wp.quat_rotate(wp.quat_inverse(q_wp_rest), rest_measure.kb_world) + joint_cable_rest_twist[j] = rest_measure.twist + + @wp.kernel def step_joint_C0_lambda( + joint_type: wp.array[int], joint_enabled: wp.array[bool], joint_parent: wp.array[int], joint_child: wp.array[int], joint_X_p: wp.array[wp.transform], joint_X_c: wp.array[wp.transform], + joint_cable_rest_kb_local: wp.array[wp.vec3], + joint_cable_rest_twist: wp.array[float], body_q_prev: wp.array[wp.transform], body_q_rest: wp.array[wp.transform], joint_constraint_start: wp.array[wp.int32], @@ -2294,9 +2923,11 @@ def step_joint_C0_lambda( ): """Per-step joint AVBD maintenance: k decay + C0 snapshot + lambda decay. - Sole owner of all joint decay. Runs every step. + Sole owner of joint decay. Cable stretch/shear and bend/twist share linear + and angular AL state blocks; non-cable drive/limit slots stay soft. """ j = wp.tid() + zero = wp.vec3(0.0) c_start = int(joint_constraint_start[j]) c_dim = int(joint_constraint_dim[j]) @@ -2309,18 +2940,36 @@ def step_joint_C0_lambda( child = joint_child[j] if not joint_enabled[j] or c_dim == 0 or child < 0: - joint_C0_lin[j] = wp.vec3(0.0) - joint_C0_ang[j] = wp.vec3(0.0) - joint_lambda_lin[j] = wp.vec3(0.0) - joint_lambda_ang[j] = wp.vec3(0.0) + joint_C0_lin[j] = zero + joint_lambda_lin[j] = zero + joint_C0_ang[j] = zero + joint_lambda_ang[j] = zero return - lin_hard = joint_is_hard[c_start] - ang_hard = 0 - if c_dim > 1: - ang_hard = joint_is_hard[c_start + 1] + jt = joint_type[j] + + # Cable has four structural slots, but AL state is stored as two vec3 + # blocks: linear = stretch/shear, angular = bend/twist. + if jt == JointType.CABLE: + stretch_idx = c_start + shear_idx = c_start + 1 + bend_idx = c_start + 2 + twist_idx = c_start + 3 + + has_linear_hard = int(0) + has_angular_hard = int(0) + if joint_is_hard[stretch_idx] == 1 or joint_is_hard[shear_idx] == 1: + has_linear_hard = 1 + if joint_is_hard[bend_idx] == 1 or joint_is_hard[twist_idx] == 1: + has_angular_hard = 1 + + if has_linear_hard == 0 and has_angular_hard == 0: + joint_C0_lin[j] = zero + joint_lambda_lin[j] = zero + joint_C0_ang[j] = zero + joint_lambda_ang[j] = zero + return - if lin_hard == 1 or ang_hard == 1: parent = joint_parent[j] if parent >= 0: X_wp = body_q_prev[parent] * joint_X_p[j] @@ -2328,35 +2977,76 @@ def step_joint_C0_lambda( X_wp = joint_X_p[j] X_wc = body_q_prev[child] * joint_X_c[j] - if lin_hard == 1: + if has_linear_hard == 1: x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) - joint_C0_lin[j] = x_c - x_p + # Store the parent-material residual [shear_x, shear_y, stretch_z]. + joint_C0_lin[j] = wp.quat_rotate_inv(wp.transform_get_rotation(X_wp), x_c - x_p) joint_lambda_lin[j] = joint_lambda_lin[j] * lambda_decay else: - joint_C0_lin[j] = wp.vec3(0.0) - joint_lambda_lin[j] = wp.vec3(0.0) + joint_C0_lin[j] = zero + joint_lambda_lin[j] = zero - if ang_hard == 1: - if parent >= 0: - X_wp_rest = body_q_rest[parent] * joint_X_p[j] - else: - X_wp_rest = joint_X_p[j] - X_wc_rest = body_q_rest[child] * joint_X_c[j] + if has_angular_hard == 1: q_wp = wp.transform_get_rotation(X_wp) q_wc = wp.transform_get_rotation(X_wc) - q_wp_rest = wp.transform_get_rotation(X_wp_rest) - q_wc_rest = wp.transform_get_rotation(X_wc_rest) - joint_C0_ang[j] = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + joint_C0_ang[j] = compute_geometric_cable_kappa_cached_z( + q_wp, + q_wc, + joint_cable_rest_kb_local[j], + joint_cable_rest_twist[j], + ) joint_lambda_ang[j] = joint_lambda_ang[j] * lambda_decay else: - joint_C0_ang[j] = wp.vec3(0.0) - joint_lambda_ang[j] = wp.vec3(0.0) + joint_C0_ang[j] = zero + joint_lambda_ang[j] = zero + return + + # Non-cable joints have at most two structural hard slots here: linear and + # angular. Drive/limit slots are always soft and ignored by this snapshot. + has_linear_hard = int(joint_is_hard[c_start]) + has_angular_hard = int(0) + if c_dim > 1 and joint_is_hard[c_start + 1] == 1: + has_angular_hard = 1 + + if has_linear_hard == 0 and has_angular_hard == 0: + joint_C0_lin[j] = zero + joint_lambda_lin[j] = zero + joint_C0_ang[j] = zero + joint_lambda_ang[j] = zero + return + + parent = joint_parent[j] + if parent >= 0: + X_wp = body_q_prev[parent] * joint_X_p[j] + else: + X_wp = joint_X_p[j] + X_wc = body_q_prev[child] * joint_X_c[j] + + if has_linear_hard == 1: + x_p = wp.transform_get_translation(X_wp) + x_c = wp.transform_get_translation(X_wc) + joint_C0_lin[j] = x_c - x_p + joint_lambda_lin[j] = joint_lambda_lin[j] * lambda_decay + else: + joint_C0_lin[j] = zero + joint_lambda_lin[j] = zero + + if has_angular_hard == 1: + q_wp = wp.transform_get_rotation(X_wp) + q_wc = wp.transform_get_rotation(X_wc) + if parent >= 0: + X_wp_rest = body_q_rest[parent] * joint_X_p[j] + else: + X_wp_rest = joint_X_p[j] + X_wc_rest = body_q_rest[child] * joint_X_c[j] + q_wp_rest = wp.transform_get_rotation(X_wp_rest) + q_wc_rest = wp.transform_get_rotation(X_wc_rest) + joint_C0_ang[j] = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + joint_lambda_ang[j] = joint_lambda_ang[j] * lambda_decay else: - joint_C0_lin[j] = wp.vec3(0.0) - joint_C0_ang[j] = wp.vec3(0.0) - joint_lambda_lin[j] = wp.vec3(0.0) - joint_lambda_ang[j] = wp.vec3(0.0) + joint_C0_ang[j] = zero + joint_lambda_ang[j] = zero @wp.kernel @@ -2427,11 +3117,6 @@ def init_body_body_contacts_avbd( body_world: wp.array[wp.int32], # Scalar parameters k_start: float, - # In/out: replayed only for matched hard contacts that were sticking. - rigid_contact_point0: wp.array[wp.vec3], - rigid_contact_point1: wp.array[wp.vec3], - rigid_contact_offset0: wp.array[wp.vec3], - rigid_contact_offset1: wp.array[wp.vec3], # Outputs contact_penalty_k: wp.array[float], contact_lambda: wp.array[wp.vec3], @@ -2441,13 +3126,13 @@ def init_body_body_contacts_avbd( ): """Restore body-body contact state from match indices. - For hard contacts: restores lambda (rotated from old to new contact frame), - penalty_k, and stick-anchor points when the previous matched contact stuck. - For soft contacts: restores penalty_k only; lambda stays zero because the - soft path is penalty-only. - Sticky hard contacts may overwrite rigid_contact_point0/1 and - rigid_contact_offset0/1 in place with the previously saved contact anchors. - C0 and decay are handled by step_body_body_contact_C0_lambda. + For hard contacts, restores lambda (rotated from the previous to the current + contact frame) and penalty_k. For soft contacts, restores penalty_k only; + lambda stays zero because the soft path is penalty-only. Contact geometry is + owned entirely by the collision pipeline: ``"latest"`` matching supplies + fresh geometry and ``"sticky"`` matching replays persistent geometry before + the solver runs. C0 and decay are handled by + :func:`step_body_body_contact_C0_lambda`. match_index[i] addresses saved contact rows from the last snapshot. Negative values (-1 unmatched, -2 broken) cold-start identically. @@ -2491,16 +3176,6 @@ def init_body_body_contacts_avbd( lam_t_old = lam_hist - n_old * lam_n lam_t_new = lam_t_old - n_new * wp.dot(lam_t_old, n_new) contact_lambda[i] = n_new * lam_n + lam_t_new - - stick_flag = history.stick_flag[slot] - # Replay saved points and offsets only for contacts whose saved - # state was sticking. Point and offset must move together; the - # surface anchor is ``point + offset``. - if stick_flag == _STICK_FLAG_ANCHOR or stick_flag == _STICK_FLAG_DEADZONE: - rigid_contact_point0[i] = history.point0[slot] - rigid_contact_point1[i] = history.point1[slot] - rigid_contact_offset0[i] = history.offset0[slot] - rigid_contact_offset1[i] = history.offset1[slot] else: contact_lambda[i] = wp.vec3(0.0) else: @@ -2511,22 +3186,12 @@ def init_body_body_contacts_avbd( @wp.kernel def snapshot_body_body_contact_history( rigid_contact_count: wp.array[int], - rigid_contact_point0: wp.array[wp.vec3], - rigid_contact_point1: wp.array[wp.vec3], - rigid_contact_offset0: wp.array[wp.vec3], - rigid_contact_offset1: wp.array[wp.vec3], rigid_contact_normal: wp.array[wp.vec3], contact_lambda: wp.array[wp.vec3], - contact_stick_flag: wp.array[wp.int32], contact_penalty_k: wp.array[float], # Persistent outputs, in RigidContactHistory order prev_lambda: wp.array[wp.vec3], - prev_stick_flag: wp.array[wp.int32], prev_penalty_k: wp.array[float], - prev_point0: wp.array[wp.vec3], - prev_point1: wp.array[wp.vec3], - prev_offset0: wp.array[wp.vec3], - prev_offset1: wp.array[wp.vec3], prev_normal: wp.array[wp.vec3], ): """Snapshot converged contact state by contact row. @@ -2539,12 +3204,7 @@ def snapshot_body_body_contact_history( return prev_lambda[i] = contact_lambda[i] - prev_stick_flag[i] = contact_stick_flag[i] prev_penalty_k[i] = contact_penalty_k[i] - prev_point0[i] = rigid_contact_point0[i] - prev_point1[i] = rigid_contact_point1[i] - prev_offset0[i] = rigid_contact_offset0[i] - prev_offset1[i] = rigid_contact_offset1[i] prev_normal[i] = rigid_contact_normal[i] @@ -2661,6 +3321,68 @@ def init_body_particle_contacts( body_particle_contact_penalty_k[i] = k_floor +@wp.func +def _cable_dahl_active_stiffness( + c_start: int, + joint_penalty_k: wp.array[float], + joint_is_hard: wp.array[wp.int32], +) -> wp.vec3: + """Current stiffness for soft cable bend/twist modes; hard modes return zero.""" + bend_idx = c_start + 2 + twist_idx = c_start + 3 + + k_bend = float(0.0) + k_bend_active = joint_penalty_k[bend_idx] + if joint_is_hard[bend_idx] == 0 and k_bend_active > 0.0: + k_bend = k_bend_active + + k_twist = float(0.0) + k_twist_active = joint_penalty_k[twist_idx] + if joint_is_hard[twist_idx] == 0 and k_twist_active > 0.0: + k_twist = k_twist_active + + return wp.vec3(k_bend, k_bend, k_twist) + + +@wp.func +def _dahl_axis_direction(d_kappa: float, d_kappa_prev: float) -> float: + """Loading direction for one scalar Dahl component.""" + direction = float(1.0) + if d_kappa > _DAHL_KAPPADOT_DEADBAND: + direction = 1.0 + elif d_kappa < -_DAHL_KAPPADOT_DEADBAND: + direction = -1.0 + else: + direction = 1.0 if d_kappa_prev >= 0.0 else -1.0 + return direction + + +@wp.func +def _advance_dahl_axis( + d_kappa: float, + d_kappa_prev: float, + sigma_prev: float, + sigma_max: float, + tau: float, +): + """Advance one scalar Dahl component from a supplied strain increment. + + Args: + d_kappa: Current strain increment. + d_kappa_prev: Previous increment used inside the direction deadband. + sigma_prev: Previous Dahl stress. + sigma_max: Magnitude of the Dahl stress envelope. + tau: Dahl transition length in strain units. + + Returns: + Updated Dahl stress and loading direction. + """ + direction = _dahl_axis_direction(d_kappa, d_kappa_prev) + exp_term = wp.exp(-direction * d_kappa / tau) + sigma = direction * sigma_max * (1.0 - exp_term) + sigma_prev * exp_term + return sigma, direction + + @wp.kernel def compute_cable_dahl_parameters( # Inputs @@ -2674,8 +3396,10 @@ def compute_cable_dahl_parameters( joint_X_c: wp.array[wp.transform], joint_constraint_start: wp.array[int], joint_penalty_k: wp.array[float], + joint_is_hard: wp.array[wp.int32], + joint_cable_rest_kb_local: wp.array[wp.vec3], + joint_cable_rest_twist: wp.array[float], body_q: wp.array[wp.transform], - body_q_rest: wp.array[wp.transform], joint_sigma_prev: wp.array[wp.vec3], joint_kappa_prev: wp.array[wp.vec3], joint_dkappa_prev: wp.array[wp.vec3], @@ -2685,62 +3409,73 @@ def compute_cable_dahl_parameters( joint_sigma_start: wp.array[wp.vec3], joint_C_fric: wp.array[wp.vec3], ): - """Compute per-step Dahl friction parameters for cable bending. - - ``joint_sigma_start`` is the linearized friction stress at step start; - ``joint_C_fric`` is d(sigma) / d(kappa). On a selected first/reset step, - curvature is based on the current start-of-step pose and the stored stress - and curvature increment are cleared. This pre-solve rebaseline covers enabled - cables; a disabled cable refreshes its history in ``update_cable_dahl_state`` - (the end-of-step finalizer) instead, so a reset while disabled is applied there. + """ + Compute shared cable Dahl hysteresis parameters (sigma0, C_fric) from + the current bend/twist strain and the stored previous Dahl state. + + The outputs are: + - sigma0: linearized friction stress at the start of the step (per component) + - C_fric: tangent stiffness d(sigma)/d(kappa) (per component) + + Dahl eps_max/tau remain per-joint scalars for compatibility with main's + custom attributes. Bend and twist still get separate envelopes through live + active stiffness. Hard or inactive subspaces produce zero Dahl stress and + tangent stiffness. + + On a selected first/reset step, curvature is rebased to the current + start-of-step pose and the stored stress and curvature increment are + cleared. This pre-solve rebaseline covers enabled cables; a disabled cable + refreshes its history in ``update_cable_dahl_state`` (the end-of-step + finalizer) instead, so a reset while disabled is applied there. """ j = wp.tid() + zero = wp.vec3(0.0) - # Only cable joints own Dahl state. - if joint_type[j] != JointType.CABLE: - joint_sigma_start[j] = wp.vec3(0.0) - joint_C_fric[j] = wp.vec3(0.0) - return + # Default to no friction; the success path overwrites both outputs below. + joint_sigma_start[j] = zero + joint_C_fric[j] = zero - # Disabled cables are not solved, and the finalizer refreshes their Dahl - # history every step, so they need no begin-of-step rebaseline. - if not joint_enabled[j]: - joint_sigma_start[j] = wp.vec3(0.0) - joint_C_fric[j] = wp.vec3(0.0) + # Only cable joints own Dahl state. Disabled cables are not solved, and + # the finalizer refreshes their Dahl history every step, so they need no + # begin-of-step rebaseline. + if not joint_enabled[j] or joint_type[j] != JointType.CABLE: return parent = joint_parent[j] child = joint_child[j] - # World-parent joints are valid; child body must exist. if child < 0: - joint_sigma_start[j] = wp.vec3(0.0) - joint_C_fric[j] = wp.vec3(0.0) return rebaseline = _world_selected(joint_world[j], pose_rebaseline_mask) - # Compute joint frames in world space (current and rest only) + eps_max = joint_eps_max[j] + tau = joint_tau[j] + c_start = joint_constraint_start[j] + k_dahl = _cable_dahl_active_stiffness(c_start, joint_penalty_k, joint_is_hard) + # A gated joint (no Dahl this step) still owes state clearing on a + # rebaseline step, so stale pre-reset stress can never resurface when a + # subspace later reactivates. + dahl_active = tau > 0.0 and eps_max > 0.0 and (k_dahl[0] > 0.0 or k_dahl[1] > 0.0 or k_dahl[2] > 0.0) + if not dahl_active and not rebaseline: + return + + # Compute joint frames in world space and the current bend/twist strain. if parent >= 0: X_wp = body_q[parent] * joint_X_p[j] - X_wp_rest = body_q_rest[parent] * joint_X_p[j] else: X_wp = joint_X_p[j] - X_wp_rest = joint_X_p[j] - X_wc = body_q[child] * joint_X_c[j] - X_wc_rest = body_q_rest[child] * joint_X_c[j] - - # Extract quaternions (current and rest configurations) q_wp = wp.transform_get_rotation(X_wp) q_wc = wp.transform_get_rotation(X_wc) - q_wp_rest = wp.transform_get_rotation(X_wp_rest) - q_wc_rest = wp.transform_get_rotation(X_wc_rest) - - # Compute curvature at the start-of-step pose. - kappa_now = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + kappa_now = compute_geometric_cable_kappa_cached_z( + q_wp, + q_wc, + joint_cable_rest_kb_local[j], + joint_cable_rest_twist[j], + ) - # Read previous state (from last converged timestep) + # Previous Dahl state (from last converged timestep). kappa_prev = joint_kappa_prev[j] d_kappa_prev = joint_dkappa_prev[j] sigma_prev = joint_sigma_prev[j] @@ -2752,54 +3487,31 @@ def compute_cable_dahl_parameters( joint_dkappa_prev[j] = d_kappa_prev joint_sigma_prev[j] = sigma_prev - # Read per-joint Dahl parameters (isotropic) - eps_max = joint_eps_max[j] - tau = joint_tau[j] - - # Use the per-joint bend stiffness from the solver constraint array (constraint slot 1 for cables). - c_start = joint_constraint_start[j] - k_bend_target = joint_penalty_k[c_start + 1] - - # Friction envelope: sigma_max = k_bend_target * eps_max. - - sigma_max = k_bend_target * eps_max - if sigma_max <= 0.0 or tau <= 0.0: - joint_sigma_start[j] = wp.vec3(0.0) - joint_C_fric[j] = wp.vec3(0.0) + if not dahl_active: return - sigma_out = wp.vec3(0.0) - C_fric_out = wp.vec3(0.0) - + d_kappa = _cable_bend_twist_delta(kappa_now, kappa_prev) + sigma_out = zero + C_fric_out = zero for axis in range(3): - kappa_i = kappa_now[axis] - kappa_i_prev = kappa_prev[axis] - sigma_i_prev = sigma_prev[axis] - - # Geometric curvature change - d_kappa_i = kappa_i - kappa_i_prev - - # Direction flag based primarily on geometric change, with stored Delta-kappa fallback - s_i = 1.0 - if d_kappa_i > _DAHL_KAPPADOT_DEADBAND: - s_i = 1.0 - elif d_kappa_i < -_DAHL_KAPPADOT_DEADBAND: - s_i = -1.0 - else: - # Within deadband: maintain previous direction from stored Delta kappa - s_i = 1.0 if d_kappa_prev[axis] >= 0.0 else -1.0 - exp_term = wp.exp(-s_i * d_kappa_i / tau) - sigma0_i = s_i * sigma_max * (1.0 - exp_term) + sigma_i_prev * exp_term - sigma0_i = wp.clamp(sigma0_i, -sigma_max, sigma_max) + sigma_max = k_dahl[axis] * eps_max + if sigma_max <= 0.0: + continue - numerator = sigma_max - s_i * sigma0_i - # Use geometric curvature change for the length scale - denominator = tau + wp.abs(d_kappa_i) + sigma0_i, direction = _advance_dahl_axis( + d_kappa[axis], + d_kappa_prev[axis], + sigma_prev[axis], + sigma_max, + tau, + ) + sigma0_i = wp.clamp(sigma0_i, -sigma_max, sigma_max) - # Store pure stiffness K = numerator / (tau + |d_kappa|) - C_fric_i = wp.max(numerator / denominator, 0.0) + # Tangent stiffness K = (sigma_max - dir*sigma0) / (tau + |d_kappa|). + numerator = sigma_max - direction * sigma0_i + denominator = tau + wp.abs(d_kappa[axis]) sigma_out[axis] = sigma0_i - C_fric_out[axis] = C_fric_i + C_fric_out[axis] = wp.max(numerator / denominator, 0.0) joint_sigma_start[j] = sigma_out joint_C_fric[j] = C_fric_out @@ -3328,6 +4040,8 @@ def solve_rigid_body( joint_X_p: wp.array[wp.transform], joint_X_c: wp.array[wp.transform], joint_axis: wp.array[wp.vec3], + joint_cable_rest_kb_local: wp.array[wp.vec3], + joint_cable_rest_twist: wp.array[float], joint_qd_start: wp.array[int], joint_target_q_start: wp.array[int], joint_constraint_start: wp.array[int], @@ -3500,6 +4214,8 @@ def solve_rigid_body( joint_X_p, joint_X_c, joint_axis, + joint_cable_rest_kb_local, + joint_cable_rest_twist, joint_qd_start, joint_target_q_start, joint_constraint_start, @@ -3579,6 +4295,8 @@ def update_duals_joint( joint_X_p: wp.array[wp.transform], joint_X_c: wp.array[wp.transform], joint_axis: wp.array[wp.vec3], + joint_cable_rest_kb_local: wp.array[wp.vec3], + joint_cable_rest_twist: wp.array[float], joint_qd_start: wp.array[int], joint_target_q_start: wp.array[int], joint_constraint_start: wp.array[int], @@ -3606,8 +4324,8 @@ def update_duals_joint( """ Update augmented-Lagrangian duals for joint constraints (per-iteration). - Structural slots (linear, angular) update lambda via ALM and ramp k, - both unconditionally. Drive/limit slots ramp k only (no lambda); + Hard structural slots update lambda via ALM; all structural slots ramp k. + Drive/limit slots ramp k only (no lambda); k is capped at ``joint_penalty_k_max`` while the force kernel applies the mode-specific stiffness cap (``min(avbd_ke, model_ke)``). """ @@ -3640,57 +4358,110 @@ def update_duals_joint( # Compute joint frames in world space if parent >= 0: X_wp = body_q[parent] * joint_X_p[j] - X_wp_rest = body_q_rest[parent] * joint_X_p[j] else: X_wp = joint_X_p[j] - X_wp_rest = joint_X_p[j] X_wc = body_q[child] * joint_X_c[j] - X_wc_rest = body_q_rest[child] * joint_X_c[j] - # CABLE joint: isotropic stretch + isotropic bend penalties (2 scalars). + # CABLE joint: fixed stretch/shear/bend/twist slots. if jt == JointType.CABLE: q_wp = wp.transform_get_rotation(X_wp) q_wc = wp.transform_get_rotation(X_wc) - q_wp_rest = wp.transform_get_rotation(X_wp_rest) - q_wc_rest = wp.transform_get_rotation(X_wc_rest) x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) - C_vec_stretch = x_c - x_p + C_vec = x_c - x_p - kappa = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + kappa = compute_geometric_cable_kappa_cached_z( + q_wp, + q_wc, + joint_cable_rest_kb_local[j], + joint_cable_rest_twist[j], + ) - # Stretch penalty update (constraint slot 0) + # Linear penalty update in the parent-material frame: local + # u = [shear_x, shear_y, stretch_z], so stretch is z and shear is xy. stretch_idx = c_start - lam_new = _update_dual_vec3( - C_vec_stretch, - joint_C0_lin[j], + shear_idx = c_start + 1 + u = wp.quat_rotate_inv(q_wp, C_vec) + lambda_lin = joint_lambda_lin[j] + C0_lin = joint_C0_lin[j] + + u_stretch = wp.vec3(0.0, 0.0, u[2]) + lam_stretch = _update_dual_vec3( + u_stretch, + wp.vec3(0.0, 0.0, C0_lin[2]), avbd_alpha, joint_penalty_k[stretch_idx], - joint_lambda_lin[j], + wp.vec3(0.0, 0.0, lambda_lin[2]), joint_is_hard[stretch_idx], ) - joint_lambda_lin[j] = lam_new + # Soft slots use pure penalty (no ALM); discard lambda so joint_lambda_* only + # carries hard-slot contributions and soft slots don't accumulate stale duals. + if joint_is_hard[stretch_idx] == 0: + lam_stretch = wp.vec3(0.0) joint_penalty_k[stretch_idx] = wp.min( - joint_penalty_k_max[stretch_idx], joint_penalty_k[stretch_idx] + beta_lin * wp.length(C_vec_stretch) + joint_penalty_k_max[stretch_idx], joint_penalty_k[stretch_idx] + beta_lin * wp.abs(u[2]) ) - # Bend penalty update (constraint slot 1) - bend_idx = c_start + 1 - lam_new = _update_dual_vec3( - kappa, - joint_C0_ang[j], + u_shear = wp.vec3(u[0], u[1], 0.0) + lam_shear = _update_dual_vec3( + u_shear, + wp.vec3(C0_lin[0], C0_lin[1], 0.0), + avbd_alpha, + joint_penalty_k[shear_idx], + wp.vec3(lambda_lin[0], lambda_lin[1], 0.0), + joint_is_hard[shear_idx], + ) + if joint_is_hard[shear_idx] == 0: + lam_shear = wp.vec3(0.0) + joint_lambda_lin[j] = lam_stretch + lam_shear + joint_penalty_k[shear_idx] = wp.min( + joint_penalty_k_max[shear_idx], joint_penalty_k[shear_idx] + beta_lin * wp.length(u_shear) + ) + + # Bend penalty update (first angular constraint slot) + bend_idx = c_start + 2 + lambda_ang = joint_lambda_ang[j] + C0_ang = joint_C0_ang[j] + kappa_bend = wp.vec3(kappa[0], kappa[1], 0.0) + lam_bend = _update_dual_vec3( + kappa_bend, + wp.vec3(C0_ang[0], C0_ang[1], 0.0), avbd_alpha, joint_penalty_k[bend_idx], - joint_lambda_ang[j], + wp.vec3(lambda_ang[0], lambda_ang[1], 0.0), joint_is_hard[bend_idx], ) - joint_lambda_ang[j] = lam_new + if joint_is_hard[bend_idx] == 0: + lam_bend = wp.vec3(0.0) joint_penalty_k[bend_idx] = wp.min( - joint_penalty_k_max[bend_idx], joint_penalty_k[bend_idx] + beta_ang * wp.length(kappa) + joint_penalty_k_max[bend_idx], joint_penalty_k[bend_idx] + beta_ang * wp.length(kappa_bend) + ) + + twist_idx = c_start + 3 + kappa_twist = wp.vec3(0.0, 0.0, kappa[2]) + lam_twist = _update_dual_vec3( + kappa_twist, + wp.vec3(0.0, 0.0, C0_ang[2]), + avbd_alpha, + joint_penalty_k[twist_idx], + wp.vec3(0.0, 0.0, lambda_ang[2]), + joint_is_hard[twist_idx], + ) + if joint_is_hard[twist_idx] == 0: + lam_twist = wp.vec3(0.0) + joint_lambda_ang[j] = lam_bend + lam_twist + joint_penalty_k[twist_idx] = wp.min( + joint_penalty_k_max[twist_idx], joint_penalty_k[twist_idx] + beta_ang * wp.length(kappa_twist) ) return + if parent >= 0: + X_wp_rest = body_q_rest[parent] * joint_X_p[j] + else: + X_wp_rest = joint_X_p[j] + X_wc_rest = body_q_rest[child] * joint_X_c[j] + # BALL joint: update isotropic linear anchor-coincidence penalty (single scalar). if jt == JointType.BALL: x_p = wp.transform_get_translation(X_wp) @@ -3999,16 +4770,12 @@ def update_duals_body_body_contacts( contact_material_mu: wp.array[float], contact_C0: wp.array[wp.vec3], avbd_alpha: float, - stick_motion_eps: float, hard_contacts: int, - body_inv_mass: wp.array[float], contact_material_ke: wp.array[float], beta: float, # Input/output contact_penalty_k: wp.array[float], contact_lambda: wp.array[wp.vec3], - # Output - contact_stick_flag: wp.array[wp.int32], ): """ Update AVBD augmented-Lagrangian duals for contact constraints (per-iteration). @@ -4085,24 +4852,6 @@ def update_duals_body_body_contacts( lam_t_new = lam_t_new * (cone_limit / lam_t_len) contact_lambda[idx] = n * lam_n_new + lam_t_new - has_kinematic = int(0) - if body_id_0 < 0 or body_id_1 < 0: - has_kinematic = int(1) - elif body_id_0 >= 0 and body_inv_mass[body_id_0] == 0.0: - has_kinematic = int(1) - elif body_id_1 >= 0 and body_inv_mass[body_id_1] == 0.0: - has_kinematic = int(1) - - flag = int(0) - if lam_n_new > 0.0 and lam_t_len <= cone_limit and wp.length(tangent_residual) < stick_motion_eps: - if has_kinematic == 1: - flag = _STICK_FLAG_ANCHOR - else: - flag = _STICK_FLAG_DEADZONE - contact_stick_flag[idx] = flag - else: - contact_stick_flag[idx] = int(0) - C_n = -contact_surface_separation(p0_world, p1_world, n, rigid_contact_margin0[idx], rigid_contact_margin1[idx]) if C_n > 0.0: contact_penalty_k[idx] = wp.min(contact_material_ke[idx], contact_penalty_k[idx] + beta * C_n) @@ -4186,13 +4935,6 @@ def update_body_velocity( dt: float, body_q: wp.array[wp.transform], body_com: wp.array[wp.vec3], - body_contact_buffer_pre_alloc: int, - body_contact_counts: wp.array[wp.int32], - body_contact_indices: wp.array[wp.int32], - contact_stick_flag: wp.array[wp.int32], - apply_stick_deadzone: int, - stick_freeze_translation_eps: float, - stick_freeze_angular_eps: float, body_q_prev: wp.array[wp.transform], body_qd: wp.array[wp.spatial_vector], body_qd_mirror: wp.array[wp.spatial_vector], @@ -4201,8 +4943,6 @@ def update_body_velocity( """ Update body velocities from position changes (world frame). - Optionally applies a tiny body-level stick-contact deadzone before - finite-difference velocity computation. Computes linear and angular velocities using finite differences. Also transfers the final body poses to body_q_out (fused copy from the in-place Gauss-Seidel iteration buffer to state_out). @@ -4214,15 +4954,6 @@ def update_body_velocity( dt: Time step. body_q: Current body transforms (world), from state_in (in-place iteration buffer). body_com: Center of mass offsets (local frame). - body_contact_buffer_pre_alloc: Per-body contact-list capacity. - body_contact_counts: Number of body-body contacts adjacent to each body. - body_contact_indices: Flat per-body contact index lists. - contact_stick_flag: Per-contact flag (0=none, ANCHOR=sticking kinematic/static, - DEADZONE=sticking dynamic-dynamic). - apply_stick_deadzone: If nonzero, enable anti-creep deadzone for bodies whose - contacts carry DEADZONE but not ANCHOR. - stick_freeze_translation_eps: Translation deadzone [m] for anti-creep snapping. - stick_freeze_angular_eps: Angular deadzone [rad] for anti-creep snapping. body_q_prev: Previous body transforms (input/output), advanced to the current pose for the next step. ``SolverVBD.reset()`` is the supported way to establish a new baseline after a discontinuous pose change. @@ -4243,27 +4974,6 @@ def update_body_velocity( q = wp.transform_get_rotation(pose) q_prev = wp.transform_get_rotation(pose_prev) - if apply_stick_deadzone != 0: - count = wp.min(body_contact_counts[tid], body_contact_buffer_pre_alloc) - offset = tid * body_contact_buffer_pre_alloc - has_anchor = int(0) - has_deadzone = int(0) - for i in range(count): - contact_idx = body_contact_indices[offset + i] - f = contact_stick_flag[contact_idx] - if f == _STICK_FLAG_ANCHOR: - has_anchor = int(1) - elif f == _STICK_FLAG_DEADZONE: - has_deadzone = int(1) - - if has_deadzone != 0 and has_anchor == 0: - translation_delta = wp.length(x - x_prev) - angular_delta = wp.length(quat_velocity(q, q_prev, 1.0)) # dt=1 gives angular displacement [rad] - if translation_delta < stick_freeze_translation_eps and angular_delta < stick_freeze_angular_eps: - pose = pose_prev - x = x_prev - q = q_prev - # Compute COM positions com_local = body_com[tid] x_com = x + wp.quat_rotate(q, com_local) @@ -4298,9 +5008,10 @@ def update_cable_dahl_state( joint_constraint_start: wp.array[int], joint_penalty_k: wp.array[float], joint_is_hard: wp.array[wp.int32], + joint_cable_rest_kb_local: wp.array[wp.vec3], + joint_cable_rest_twist: wp.array[float], # Body states (final, after solver convergence) body_q: wp.array[wp.transform], - body_q_rest: wp.array[wp.transform], # Dahl model parameters (PER-JOINT arrays, isotropic) joint_eps_max: wp.array[float], joint_tau: wp.array[float], @@ -4310,121 +5021,83 @@ def update_cable_dahl_state( joint_dkappa_prev: wp.array[wp.vec3], # input/output (stores Delta kappa) ): """ - Post-iteration kernel: update Dahl hysteresis state after solver convergence (component-wise). - - Stores final curvature, friction stress, and curvature Delta kappa for the next step. Each - curvature component (x, y, z) is updated independently to preserve path-dependent memory. + Persist cable Dahl hysteresis state after solver convergence. - Args: - joint_type: Joint type (only updates for cable joints) - joint_parent, joint_child: Parent/child body indices - joint_X_p, joint_X_c: Joint frames in parent/child - joint_constraint_start: Start index per joint in the solver constraint layout - joint_penalty_k: Per-constraint penalty stiffness; for cables, bend slot stores effective per-joint bend stiffness [N*m] - body_q: Final body transforms (after convergence) - body_q_rest: Rest body transforms - joint_sigma_prev: Friction stress state (read old, write new), wp.vec3 per joint - joint_kappa_prev: Curvature state (read old, write new), wp.vec3 per joint - joint_dkappa_prev: Delta-kappa state (write new), wp.vec3 per joint - joint_eps_max: Maximum persistent strain [rad] (scalar per joint) - joint_tau: Memory decay length [rad] (scalar per joint) + State is diagonal in [bend_x, bend_y, twist_z]. Only soft modes with active + stiffness are advanced; inactive modes clear stress and use final strain as + the next baseline. """ j = wp.tid() + zero = wp.vec3(0.0) - # Only update cable joints if joint_type[j] != JointType.CABLE: return - # Get parent and child body indices parent = joint_parent[j] child = joint_child[j] - - # World-parent joints are valid; child body must exist. if child < 0: return - # Compute joint frames in world space (final state) if parent >= 0: X_wp = body_q[parent] * joint_X_p[j] - X_wp_rest = body_q_rest[parent] * joint_X_p[j] else: X_wp = joint_X_p[j] - X_wp_rest = joint_X_p[j] X_wc = body_q[child] * joint_X_c[j] - X_wc_rest = body_q_rest[child] * joint_X_c[j] q_wp = wp.transform_get_rotation(X_wp) q_wc = wp.transform_get_rotation(X_wc) - q_wp_rest = wp.transform_get_rotation(X_wp_rest) - q_wc_rest = wp.transform_get_rotation(X_wc_rest) - # Compute final curvature vector at end of timestep - kappa_final = compute_kappa(q_wp, q_wc, q_wp_rest, q_wc_rest) + kappa_final = compute_geometric_cable_kappa_cached_z( + q_wp, + q_wc, + joint_cable_rest_kb_local[j], + joint_cable_rest_twist[j], + ) - # Refresh Dahl state so toggling enabled/hard does not see stale values. - c_start_dahl = joint_constraint_start[j] - if not joint_enabled[j] or joint_is_hard[c_start_dahl + 1] == 1: + c_start = joint_constraint_start[j] + k_dahl = _cable_dahl_active_stiffness(c_start, joint_penalty_k, joint_is_hard) + + # Inactive modes clear stress and use the current strain as the next baseline. + if not joint_enabled[j] or (k_dahl[0] <= 0.0 and k_dahl[1] <= 0.0 and k_dahl[2] <= 0.0): joint_kappa_prev[j] = kappa_final - joint_sigma_prev[j] = wp.vec3(0.0) - joint_dkappa_prev[j] = wp.vec3(0.0) + joint_sigma_prev[j] = zero + joint_dkappa_prev[j] = zero return - # Read stored Dahl state (component-wise vectors) - kappa_old = joint_kappa_prev[j] # stored curvature - d_kappa_old = joint_dkappa_prev[j] # stored Delta kappa - sigma_old = joint_sigma_prev[j] # stored friction stress + # Stored Dahl state from the previous converged timestep. + kappa_old = joint_kappa_prev[j] + d_kappa_old = joint_dkappa_prev[j] + sigma_old = joint_sigma_prev[j] + d_kappa = _cable_bend_twist_delta(kappa_final, kappa_old) - # Read per-joint Dahl parameters (isotropic) eps_max = joint_eps_max[j] # Maximum persistent strain [rad] tau = joint_tau[j] # Memory decay length [rad] - # Bend stiffness is stored in constraint slot 1 for cable joints. - c_start = joint_constraint_start[j] - k_bend_target = joint_penalty_k[c_start + 1] # [N*m] - - # Friction envelope: sigma_max = k_bend_target * eps_max. - sigma_max = k_bend_target * eps_max # [N*m] - - # Early-out: disable friction if envelope is zero/invalid - if sigma_max <= 0.0 or tau <= 0.0: - joint_sigma_prev[j] = wp.vec3(0.0) + if eps_max <= 0.0 or tau <= 0.0: + joint_sigma_prev[j] = zero joint_kappa_prev[j] = kappa_final - joint_dkappa_prev[j] = kappa_final - kappa_old # store Delta kappa + joint_dkappa_prev[j] = d_kappa return - # Update each component independently (3 separate hysteresis loops) - sigma_final_out = wp.vec3(0.0) - d_kappa_out = wp.vec3(0.0) + sigma_final_out = zero + d_kappa_out = zero for axis in range(3): - # Get component values - kappa_i_final = kappa_final[axis] - kappa_i_prev = kappa_old[axis] - d_kappa_i_prev = d_kappa_old[axis] - sigma_i_prev = sigma_old[axis] - - # Curvature change for this component - d_kappa_i = kappa_i_final - kappa_i_prev - - # Direction flag (same logic as pre-iteration kernel), in kappa-space - s_i = 1.0 - if d_kappa_i > _DAHL_KAPPADOT_DEADBAND: - s_i = 1.0 - elif d_kappa_i < -_DAHL_KAPPADOT_DEADBAND: - s_i = -1.0 - else: - # Within deadband: maintain previous direction - s_i = 1.0 if d_kappa_i_prev >= 0.0 else -1.0 - - # sigma_i_next = s_i*sigma_max * [1 - exp(-s_i*d_kappa_i/tau)] + sigma_i_prev * exp(-s_i*d_kappa_i/tau) - exp_term = wp.exp(-s_i * d_kappa_i / tau) - sigma_i_next = s_i * sigma_max * (1.0 - exp_term) + sigma_i_prev * exp_term + sigma_max = k_dahl[axis] * eps_max # [N*m] + if sigma_max <= 0.0: + continue - # Store component results + sigma_i_next, _direction = _advance_dahl_axis( + d_kappa[axis], + d_kappa_old[axis], + sigma_old[axis], + sigma_max, + tau, + ) sigma_final_out[axis] = sigma_i_next - d_kappa_out[axis] = d_kappa_i + d_kappa_out[axis] = d_kappa[axis] - # Store final vector state for next timestep + # Store final vector state for next timestep: [bend_x, bend_y, twist_z]. joint_sigma_prev[j] = sigma_final_out joint_kappa_prev[j] = kappa_final joint_dkappa_prev[j] = d_kappa_out diff --git a/newton/_src/solvers/vbd/solver_vbd.py b/newton/_src/solvers/vbd/solver_vbd.py index 55b15fa1be..c5c25fb07d 100644 --- a/newton/_src/solvers/vbd/solver_vbd.py +++ b/newton/_src/solvers/vbd/solver_vbd.py @@ -22,6 +22,8 @@ State, StateFlags, ) +from ...sim.collide import _count_soft_particle_rigid_contact_pairs +from ...utils import is_graph_capture_allocation_enabled from ...utils.deprecation import deprecate_nonkeyword_arguments from ..coupled.interface import CouplingInterface from ..solver import SolverBase @@ -61,6 +63,7 @@ init_body_body_contact_materials, init_body_body_contacts_avbd, init_body_particle_contacts, + init_cable_rest_bend_twist, reset_rigid_state, snapshot_body_body_contact_history, solve_rigid_body, @@ -104,18 +107,21 @@ class SolverVBD(SolverBase, CouplingInterface): use augmented-Lagrangian state. Non-cable structural joint slots default to **hard mode** (augmented Lagrangian - with persistent lambda and C0 stabilization). Cable stretch and bend default to - **soft mode**. Joint hard/soft mode is initialized from the optional - ``model.vbd.joint_is_hard`` custom attribute; author values at joint creation, - before constructing the solver. The hard/soft mode can also be changed per - slot at runtime via :meth:`set_joint_constraint_mode`. + with persistent lambda and C0 stabilization) and are initialized from the + optional ``model.vbd.joint_is_hard`` custom attribute; author values at joint + creation, before constructing the solver. Cable stretch, shear, bend, and + twist always initialize to **soft mode** regardless of ``joint_is_hard`` and + are switched only at runtime. The hard/soft mode can be changed per slot at + runtime via :meth:`set_joint_constraint_mode` (for both cable and non-cable + joints). Joint limitations: - Supported joint types: BALL, FIXED, FREE, REVOLUTE, PRISMATIC, D6, CABLE. DISTANCE joints are not supported. - :attr:`~newton.Model.joint_enabled` is supported for all joint types. - :attr:`~newton.Model.joint_target_ke`/:attr:`~newton.Model.joint_target_kd` are supported - for REVOLUTE, PRISMATIC, D6 (as drives), and CABLE (as stretch/bend stiffness and damping). + for REVOLUTE, PRISMATIC, D6 (as drives), and CABLE (as stretch, shear, + bend, and twist stiffness and damping). VBD interprets ``kd`` as absolute damping in physical units. - :attr:`~newton.Model.joint_limit_lower`/:attr:`~newton.Model.joint_limit_upper` and :attr:`~newton.Model.joint_limit_ke`/:attr:`~newton.Model.joint_limit_kd` are supported @@ -128,15 +134,22 @@ class SolverVBD(SolverBase, CouplingInterface): See :ref:`Joint feature support` for the full comparison across solvers. Buffer sizing: - SolverVBD pre-allocates contact state from capacities populated by - :class:`~newton.CollisionPipeline` when available; otherwise, the first - :meth:`step` lazily sizes buffers from ``Contacts``. During CUDA graph - recording, ordinary lazy resizing is supported only when Warp's memory pool - is enabled; otherwise, the solver raises with guidance to pre-size before - capture. Rigid contact history must be allocated before capture regardless - of memory-pool support. With ``rigid_contact_history=True``, construct - :class:`~newton.CollisionPipeline` before ``SolverVBD``, or run one - uncaptured solver step before capture. + Body-body contact state is pre-allocated from ``model.rigid_contact_max`` when a + :class:`~newton.CollisionPipeline` has already published it and this solver owns the + rigid system. Body-particle contact state is pre-sized from a world-aware + particle-shape pair count, which excludes the + ``enable_rigid_soft_full_surface_contact`` edge/face headroom. Both grow from + ``Contacts`` on the first :meth:`step`, and the rigid contact force outputs grow in + :meth:`collect_rigid_contact_forces`. During graph capture, + ordinary lazy resizing is supported on CPU and on CUDA with Warp's + stream-ordered memory pool enabled; otherwise the solver raises with + guidance to pre-size before capture. Rigid contact history is + cross-replay-persistent state, so it must always be allocated before + capture regardless of the device's allocation-during-capture support -- + allocating it inside a graph records a `wp.zeros` fill that wipes the + warm-start buffers on every replay. With ``rigid_contact_history=True``, + construct :class:`~newton.CollisionPipeline` before ``SolverVBD``, or run + one uncaptured solver step before capture. References: - Anka He Chen, Ziheng Liu, Yin Yang, and Cem Yuksel. 2024. Vertex Block Descent. ACM Trans. Graph. 43, 4, Article 116 (July 2024), 16 pages. @@ -157,6 +170,9 @@ class SolverVBD(SolverBase, CouplingInterface): transforms must match the joint angles at solver creation time (see example below). + For CUDA graph capture, the recommended construction order is + ``CollisionPipeline`` -> ``Contacts`` -> ``SolverVBD``, all before capture. + Example ------- @@ -167,14 +183,15 @@ class SolverVBD(SolverBase, CouplingInterface): model = builder.finalize() + collision_pipeline = newton.CollisionPipeline(model) + contacts = collision_pipeline.contacts() + solver = newton.solvers.SolverVBD(model) - # Initialize states and contacts + # Initialize states and control state_in = model.state() state_out = model.state() control = model.control() - collision_pipeline = newton.CollisionPipeline(model) - contacts = collision_pipeline.contacts() # Simulation loop for i in range(100): @@ -186,19 +203,25 @@ class SolverVBD(SolverBase, CouplingInterface): class JointSlot: """Named constraint slot indices for :meth:`set_joint_constraint_mode`. - The first two solver constraint slots are structural where present: - - CABLE: LINEAR/STRETCH -> stretch, ANGULAR/BEND -> bend - - BALL: LINEAR only - - FIXED/REVOLUTE/PRISMATIC/D6: LINEAR and ANGULAR + Structural constraint slots by joint type: + - CABLE: STRETCH=0, SHEAR=1, BEND=2, TWIST=3 + - BALL: LINEAR=0 only + - FIXED/REVOLUTE/PRISMATIC/D6: LINEAR=0, ANGULAR=1 - Drive/limit slots start at slot 2 and are not represented here. - STRETCH and BEND are cable-only aliases for LINEAR and ANGULAR. + STRETCH/SHEAR/BEND/TWIST are cable-only names for the SolverVBD cable + layout emitted by the builder cable APIs. Only structural slots are named + here; per-DOF drive/limit slots (slot 2+ on non-cable joints) are not. """ + # Non-cable structural slots. LINEAR = 0 ANGULAR = 1 + # Cable structural slots (all four are linear/angular cable constraints; + # they are not the non-cable LINEAR/ANGULAR despite STRETCH sharing index 0). STRETCH = 0 - BEND = 1 + SHEAR = 1 + BEND = 2 + TWIST = 3 @deprecate_nonkeyword_arguments def __init__( @@ -233,10 +256,10 @@ def __init__( rigid_avbd_gamma: float = 0.999, # Per-step decay for penalty k and persisted hard-mode lambda # Rigid body - contacts rigid_contact_hard: bool = True, # Body-body contacts: hard=AL duals+C0, soft=penalty only - rigid_contact_history: bool = False, # Body-body contact warm-start (hard: k+duals+anchors; soft: k) - rigid_contact_stick_motion_eps: float = 1.0e-4, # Sticky contact residual threshold; 0 disables point replay - rigid_contact_stick_freeze_translation_eps: float = 1.0e-4, # Deadzone snap translation threshold; 0 disables snap - rigid_contact_stick_freeze_angular_eps: float = 1.0e-4, # Deadzone snap angular threshold; 0 disables snap + rigid_contact_history: bool = False, # Body-body contact numeric warm-start (hard: k+duals; soft: k) + rigid_contact_stick_motion_eps: float | None = None, # Deprecated and ignored + rigid_contact_stick_freeze_translation_eps: float | None = None, # Deprecated and ignored + rigid_contact_stick_freeze_angular_eps: float | None = None, # Deprecated and ignored rigid_contact_k_start: float = 1.0e2, # Body-body/body-particle penalty seed when ramping is enabled rigid_body_contact_buffer_size: int = 64, # Per-body body-body contact list capacity rigid_body_particle_contact_buffer_size: int = 256, # Per-body soft-contact list capacity (particle + edge/face) @@ -320,27 +343,29 @@ def __init__( reference scheme. Lower values decay faster, improving stability at the cost of slower convergence. rigid_contact_hard: Whether body-body rigid contacts use hard mode (augmented Lagrangian with persistent lambda and C0 stabilization) or soft mode (penalty only). - rigid_contact_history: Whether to persist body-body contact state across steps using - ``Contacts.rigid_contact_match_index`` from the collision pipeline. For hard contacts, - restores lambda, penalty k, and sticky contact anchors; C0 is recomputed each step. - For soft contacts, only restored penalty k affects the solve (useful with ramping). - Requires contacts with ``rigid_contact_match_index`` populated; use - ``CollisionPipeline(contact_matching="latest")`` for VBD warm-starting. Ignored - when ``integrate_with_external_rigid_solver=True`` or ``model.body_count == 0``. - For CUDA graph capture, construct :class:`~newton.CollisionPipeline` before - ``SolverVBD`` so history is pre-allocated, or run one uncaptured solver step - before capture. - rigid_contact_stick_motion_eps: Tangential contact residual threshold for marking hard - body-body contacts as sticking. Sticking contacts may replay contact points when - ``rigid_contact_history=True``; dynamic-dynamic sticking contacts may also use the - body-level deadzone snap. Set to ``0.0`` to disable sticky flags while preserving - lambda and penalty warm-starting. - rigid_contact_stick_freeze_translation_eps: World-space translation threshold for the - body-level deadzone snap on dynamic-dynamic sticking contacts. Set to ``0.0`` to - disable translation snapping. - rigid_contact_stick_freeze_angular_eps: Angular threshold [rad] for the body-level - deadzone snap on dynamic-dynamic sticking contacts. Set to ``0.0`` to disable - angular snapping. + rigid_contact_history: Whether to persist body-body numeric contact state + across steps using ``Contacts.rigid_contact_match_index``. Hard contacts + restore lambda and penalty k; soft contacts restore penalty k only. + Contact geometry remains owned by the collision pipeline. Requires + ``CollisionPipeline(contact_matching="latest")`` or ``"sticky"``. + Ignored when ``integrate_with_external_rigid_solver=True`` or + ``model.body_count == 0``. During graph capture, construct the + collision pipeline before ``SolverVBD`` so history is pre-allocated, + or run one uncaptured solver step before capture. + rigid_contact_stick_motion_eps: Deprecated and ignored. SolverVBD no longer + classifies contacts as sticking. Use + ``CollisionPipeline(contact_matching="sticky", + contact_matching_pos_threshold=...)`` for persistent contact geometry. + + .. deprecated:: 1.5 + rigid_contact_stick_freeze_translation_eps: Deprecated and ignored. The + SolverVBD body-level contact deadzone was removed. + + .. deprecated:: 1.5 + rigid_contact_stick_freeze_angular_eps: Deprecated and ignored. The + SolverVBD body-level contact deadzone was removed. + + .. deprecated:: 1.5 rigid_contact_k_start: Body-body and body-particle contact penalty seed for AVBD ramping. Used when ``rigid_avbd_linear_beta`` (or ``rigid_avbd_beta`` fallback) is greater than zero. When the linear beta is 0, k is fixed at the contact stiffness regardless of this value. @@ -374,7 +399,7 @@ def __init__( Setting them too small may result in undetected collisions (particles) or contact overflow (rigid body contacts). Setting them excessively large may increase memory usage and degrade performance. - - Dahl hysteresis friction for cable bending is controlled by custom model attributes + - Dahl hysteresis friction for cable angular response is controlled by custom model attributes ``model.vbd.dahl_eps_max`` and ``model.vbd.dahl_tau``. Register them with ``SolverVBD.register_custom_attributes`` before building the model. Dahl friction is enabled only when positive Dahl parameters are authored. @@ -384,6 +409,23 @@ def __init__( raise ValueError(f"rigid_avbd_beta must be >= 0, got {rigid_avbd_beta}") rigid_avbd_linear_beta = rigid_avbd_linear_beta if rigid_avbd_linear_beta is not None else rigid_avbd_beta rigid_avbd_angular_beta = rigid_avbd_angular_beta if rigid_avbd_angular_beta is not None else rigid_avbd_beta + if ( + rigid_contact_stick_motion_eps is not None + or rigid_contact_stick_freeze_translation_eps is not None + or rigid_contact_stick_freeze_angular_eps is not None + ): + warnings.warn( + "SolverVBD rigid_contact_stick_motion_eps, " + "rigid_contact_stick_freeze_translation_eps, and " + "rigid_contact_stick_freeze_angular_eps are deprecated and ignored, " + "and will be removed in a future release. " + "Use CollisionPipeline(contact_matching='sticky', " + "contact_matching_pos_threshold=...) for persistent contact geometry. " + "The SolverVBD body-level contact deadzone was removed.", + DeprecationWarning, + # __init__ is wrapped by @deprecate_nonkeyword_arguments. + stacklevel=3, + ) super().__init__(model) @@ -466,9 +508,6 @@ def __init__( rigid_avbd_contact_alpha, rigid_contact_hard, rigid_contact_history, - rigid_contact_stick_motion_eps, - rigid_contact_stick_freeze_translation_eps, - rigid_contact_stick_freeze_angular_eps, rigid_contact_k_start, rigid_body_contact_buffer_size, rigid_body_particle_contact_buffer_size, @@ -594,9 +633,6 @@ def _init_rigid_system( rigid_avbd_contact_alpha: float | None, rigid_contact_hard: bool, rigid_contact_history: bool, - rigid_contact_stick_motion_eps: float, - rigid_contact_stick_freeze_translation_eps: float, - rigid_contact_stick_freeze_angular_eps: float, rigid_contact_k_start: float, rigid_body_contact_buffer_size: int, rigid_body_particle_contact_buffer_size: int, @@ -628,17 +664,6 @@ def _init_rigid_system( raise ValueError(f"rigid_avbd_gamma must be in [0, 1], got {rigid_avbd_gamma}") if rigid_contact_k_start < 0: raise ValueError(f"rigid_contact_k_start must be >= 0, got {rigid_contact_k_start}") - if rigid_contact_stick_motion_eps < 0: - raise ValueError(f"rigid_contact_stick_motion_eps must be >= 0, got {rigid_contact_stick_motion_eps}") - if rigid_contact_stick_freeze_translation_eps < 0: - raise ValueError( - "rigid_contact_stick_freeze_translation_eps must be >= 0, " - f"got {rigid_contact_stick_freeze_translation_eps}" - ) - if rigid_contact_stick_freeze_angular_eps < 0: - raise ValueError( - f"rigid_contact_stick_freeze_angular_eps must be >= 0, got {rigid_contact_stick_freeze_angular_eps}" - ) if rigid_joint_linear_k_start < 0: raise ValueError(f"rigid_joint_linear_k_start must be >= 0, got {rigid_joint_linear_k_start}") if rigid_joint_angular_k_start < 0: @@ -662,11 +687,6 @@ def _init_rigid_system( else: self.rigid_contact_alpha = rigid_avbd_alpha - self.rigid_contact_stick_motion_eps = rigid_contact_stick_motion_eps - # DEADZONE body-snap thresholds; suppressed by _STICK_FLAG_ANCHOR. - self.rigid_contact_stick_freeze_translation_eps = rigid_contact_stick_freeze_translation_eps - self.rigid_contact_stick_freeze_angular_eps = rigid_contact_stick_freeze_angular_eps - # Joint constraint stiffness and damping for non-cable structural joints self.rigid_joint_linear_ke = rigid_joint_linear_ke self.rigid_joint_angular_ke = rigid_joint_angular_ke @@ -742,16 +762,10 @@ def _init_rigid_system( self.body_body_contact_material_mu = wp.zeros(0, dtype=float, device=self.device) self.body_body_contact_lambda = wp.zeros(0, dtype=wp.vec3, device=self.device) self.body_body_contact_C0 = wp.zeros(0, dtype=wp.vec3, device=self.device) - self.body_body_contact_stick_flag = wp.zeros(0, dtype=wp.int32, device=self.device) # Rigid contact warm-start buffers. self._prev_contact_lambda = None - self._prev_contact_stick_flag = None self._prev_contact_penalty_k = None - self._prev_contact_point0 = None - self._prev_contact_point1 = None - self._prev_contact_offset0 = None - self._prev_contact_offset1 = None self._prev_contact_normal = None # Joint augmented-Lagrangian state (vec3, per-joint, bilateral) @@ -760,7 +774,7 @@ def _init_rigid_system( self.joint_C0_lin = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) self.joint_C0_ang = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) - # Dahl friction state (cable bending plasticity, persistent across timesteps) + # Dahl friction state (cable angular hysteresis, persistent across timesteps) self.joint_sigma_prev = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) self.joint_kappa_prev = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) self.joint_dkappa_prev = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) @@ -788,6 +802,14 @@ def _init_rigid_system( self.joint_dahl_tau = wp.zeros(model.joint_count, dtype=float, device=self.device) self.enable_dahl_friction = False + # Per-joint DER rest invariants, refreshed at init and on model change + # (see _refresh_cable_rest_bend_twist_cache): the parent-local rest + # curvature binormal (bend) and the rest transported-material twist. + # Split cables use local +Z as the material tangent (a SolverVBD convention). + self.joint_cable_rest_kb_local = wp.zeros(model.joint_count, dtype=wp.vec3, device=self.device) + self.joint_cable_rest_twist = wp.zeros(model.joint_count, dtype=float, device=self.device) + self._refresh_cable_rest_bend_twist_cache() + # ------------------------------------------------------------- # Body-particle interaction shared state. # ------------------------------------------------------------- @@ -798,7 +820,10 @@ def _init_rigid_system( # Zero-length body poses for static-shape contact kernels when State.body_q is absent. self._empty_body_q = wp.empty(0, dtype=wp.transform, device=self.device) if model.particle_count > 0 and model.shape_count > 0: - self._init_body_particle_contact_state(model.shape_count * model.particle_count) + # Not shape_count * particle_count: that counts cross-world pairs, so it is quadratic in + # world count and can exceed Warp's int32 array shape limit. A hint only -- the first step + # grows this to contacts.soft_contact_max, raising if capture cannot allocate. + self._init_body_particle_contact_state(_count_soft_particle_rigid_contact_pairs(model)) # Kinematic body support: create effective inv_mass / inv_inertia arrays # with kinematic bodies zeroed out. @@ -835,6 +860,8 @@ def notify_model_changed(self, flags: ModelFlags | int) -> None: self._apply_module_options() if flags & (ModelFlags.BODY_PROPERTIES | ModelFlags.BODY_INERTIAL_PROPERTIES): self._refresh_kinematic_state() + if flags & (ModelFlags.JOINT_PROPERTIES | ModelFlags.BODY_PROPERTIES): + self._refresh_cable_rest_bend_twist_cache() @override def coupling_supports_inertial_property_refresh(self) -> bool: @@ -1134,7 +1161,6 @@ def _init_body_body_contact_state(self, rigid_contact_max: int) -> None: self.body_body_contact_material_mu = wp.zeros(rigid_contact_max, dtype=float, device=self.device) self.body_body_contact_lambda = wp.zeros(rigid_contact_max, dtype=wp.vec3, device=self.device) self.body_body_contact_C0 = wp.zeros(rigid_contact_max, dtype=wp.vec3, device=self.device) - self.body_body_contact_stick_flag = wp.zeros(rigid_contact_max, dtype=wp.int32, device=self.device) def _init_body_particle_contact_state(self, soft_contact_max: int) -> None: """Allocate body-particle material arrays sized to the given soft contact capacity.""" @@ -1147,26 +1173,55 @@ def _init_rigid_contact_warmstart(self, rigid_contact_max: int) -> None: """Allocate rigid contact warm-start buffers.""" cap = max(1, rigid_contact_max) self._prev_contact_lambda = wp.zeros(cap, dtype=wp.vec3, device=self.device) - self._prev_contact_stick_flag = wp.zeros(cap, dtype=wp.int32, device=self.device) self._prev_contact_penalty_k = wp.zeros(cap, dtype=float, device=self.device) - self._prev_contact_point0 = wp.zeros(cap, dtype=wp.vec3, device=self.device) - self._prev_contact_point1 = wp.zeros(cap, dtype=wp.vec3, device=self.device) - self._prev_contact_offset0 = wp.zeros(cap, dtype=wp.vec3, device=self.device) - self._prev_contact_offset1 = wp.zeros(cap, dtype=wp.vec3, device=self.device) self._prev_contact_normal = wp.zeros(cap, dtype=wp.vec3, device=self.device) def _raise_if_capturing_resize(self, name: str, current: int, required: int) -> None: - from ...utils import is_graph_capture_allocation_enabled # noqa: PLC0415 - if self.device.is_capturing and not is_graph_capture_allocation_enabled(self.device): raise RuntimeError( f"SolverVBD {name} buffer needs to grow from {current} to {required} " "during graph capture, but allocation during capture is not enabled on this device. " - "Pre-size before capture by constructing CollisionPipeline before SolverVBD, " - "passing explicit rigid_contact_max/soft_contact_max to CollisionPipeline, or running one " - "uncaptured step/force-collection pass." + "Run one uncaptured step (or force-collection pass) before capture so the contact " + "buffers are sized for the scene, or enable Warp's stream-ordered memory pool on this device. " + "Rigid buffers can also be pre-sized by constructing CollisionPipeline before SolverVBD, " + "which publishes model.rigid_contact_max; there is no equivalent for body-particle contacts." ) + def _refresh_cable_rest_bend_twist_cache(self) -> None: + """(Re)compute cable rest bend/twist invariants from the current rest pose. + + Called once at init and again from ``notify_model_changed`` whenever joint + frames or the rest pose change. + """ + # The cache is only allocated when SolverVBD integrates the rigid system + # (see _init_rigid_system); skip when bodies are handled externally. + if self.integrate_with_external_rigid_solver: + return + if self.model.joint_count == 0 or self.model.body_count == 0: + return + + joint_type_np = self._to_numpy(self.model.joint_type, dtype=np.int32) + if not np.any(joint_type_np == int(JointType.CABLE)): + return + + wp.launch( + kernel=init_cable_rest_bend_twist, + dim=self.model.joint_count, + inputs=[ + self.model.joint_type, + self.model.joint_parent, + self.model.joint_child, + self.model.joint_X_p, + self.model.joint_X_c, + self.model.body_q, + ], + outputs=[ + self.joint_cable_rest_kb_local, + self.joint_cable_rest_twist, + ], + device=self.device, + ) + @staticmethod def _to_numpy(arr, dtype=None): """Transfer a Warp array to CPU and return as numpy, optionally casting dtype.""" @@ -1178,7 +1233,7 @@ def _init_joint_constraint_layout(self) -> None: """Initialize VBD-owned joint constraint indexing. VBD stores and adapts penalty stiffness values for scalar constraint components: - - CABLE: 2 scalars (stretch/linear, bend/angular) + - CABLE: 4 scalars (stretch, shear, bend, twist) - BALL: 1 scalar (isotropic linear anchor-coincidence) - FIXED: 2 scalars (isotropic linear + isotropic angular) - REVOLUTE: 3 scalars (isotropic linear + 2-DOF perpendicular angular + angular drive/limit) @@ -1198,7 +1253,16 @@ def _init_joint_constraint_layout(self) -> None: dim_np = np.zeros((n_j,), dtype=np.int32) for j in range(n_j): if jt[j] == JointType.CABLE: - dim_np[j] = 2 + lin_count = int(jdof_dim[j, 0]) + ang_count = int(jdof_dim[j, 1]) + if lin_count != 2 or ang_count != 2: + raise RuntimeError( + "SolverVBD rigid joints: JointType.CABLE requires the split " + "stretch/shear/bend/twist layout emitted by the cable builder APIs " + f"(got linear={lin_count}, angular={ang_count}) " + f"for joint {j}." + ) + dim_np[j] = 4 elif jt[j] == JointType.BALL: dim_np[j] = 1 elif jt[j] == JointType.FIXED: @@ -1273,7 +1337,7 @@ def _init_joint_penalty_k(self): jdof_dim = self._to_numpy(self.model.joint_dof_dim, dtype=int) jc_start = self._to_numpy(self.joint_constraint_start, dtype=np.int32) - # Per-joint hard/soft mode from model attribute (default=1, hard). + # Per-joint hard/soft mode for non-cable structural slots. vbd_attrs: Any = getattr(self.model, "vbd", None) if vbd_attrs is not None and hasattr(vbd_attrs, "joint_is_hard"): j_is_hard = self._to_numpy(vbd_attrs.joint_is_hard, dtype=np.int32) @@ -1290,21 +1354,43 @@ def _init_joint_penalty_k(self): if jt[j] == JointType.CABLE: c0 = int(jc_start[j]) dof0 = int(jdofs[j]) - if dof0 < 0 or (dof0 + 1) >= len(jtarget_ke) or (dof0 + 1) >= len(jtarget_kd): + if dof0 < 0 or (dof0 + 3) >= len(jtarget_ke) or (dof0 + 3) >= len(jtarget_kd): raise RuntimeError( - "SolverVBD _init_joint_penalty_k: JointType.CABLE requires 2 DOF entries in " + "SolverVBD _init_joint_penalty_k: JointType.CABLE requires " + "4 DOF entries in " "model.joint_target_ke/kd starting at joint_qd_start[j]. " f"Got joint_index={j}, joint_qd_start={dof0}, " f"len(joint_target_ke)={len(jtarget_ke)}, len(joint_target_kd)={len(jtarget_kd)}." ) - ke_stretch = jtarget_ke[dof0] - ke_bend = jtarget_ke[dof0 + 1] - joint_k_max_np[c0] = ke_stretch - joint_k_max_np[c0 + 1] = ke_bend - joint_k_init_np[c0] = ke_stretch if lin_k_start is None else min(lin_k_start, ke_stretch) - joint_k_init_np[c0 + 1] = ke_bend if ang_k_start is None else min(ang_k_start, ke_bend) - joint_kd_np[c0] = jtarget_kd[dof0] - joint_kd_np[c0 + 1] = jtarget_kd[dof0 + 1] + stretch_slot = c0 + shear_slot = c0 + 1 + bend_slot = c0 + 2 + twist_slot = c0 + 3 + + stretch_dof = dof0 + shear_dof = dof0 + 1 + bend_dof = dof0 + 2 + twist_dof = dof0 + 3 + + ke_stretch = jtarget_ke[stretch_dof] + ke_shear = jtarget_ke[shear_dof] + ke_bend = jtarget_ke[bend_dof] + ke_twist = jtarget_ke[twist_dof] + + joint_k_max_np[stretch_slot] = ke_stretch + joint_k_max_np[shear_slot] = ke_shear + joint_k_max_np[bend_slot] = ke_bend + joint_k_max_np[twist_slot] = ke_twist + + joint_k_init_np[stretch_slot] = ke_stretch if lin_k_start is None else min(lin_k_start, ke_stretch) + joint_k_init_np[shear_slot] = ke_shear if lin_k_start is None else min(lin_k_start, ke_shear) + joint_k_init_np[bend_slot] = ke_bend if ang_k_start is None else min(ang_k_start, ke_bend) + joint_k_init_np[twist_slot] = ke_twist if ang_k_start is None else min(ang_k_start, ke_twist) + + joint_kd_np[stretch_slot] = jtarget_kd[stretch_dof] + joint_kd_np[shear_slot] = jtarget_kd[shear_dof] + joint_kd_np[bend_slot] = jtarget_kd[bend_dof] + joint_kd_np[twist_slot] = jtarget_kd[twist_dof] elif jt[j] == JointType.BALL: c0 = int(jc_start[j]) joint_k_max_np[c0] = structural_linear_ke @@ -1454,29 +1540,36 @@ def _init_joint_rest_angle(self): @override @classmethod - def register_custom_attributes(cls, builder: ModelBuilder, *, dahl_defaults_enabled: bool = True) -> None: + def register_custom_attributes(cls, builder: ModelBuilder, *, dahl_defaults_enabled: bool = False) -> None: """Register SolverVBD custom Model attributes. Currently registers: - - ``vbd:joint_is_hard`` for per-joint hard/soft constraint mode - - ``vbd:dahl_eps_max`` and ``vbd:dahl_tau`` for optional Dahl cable friction + - ``vbd:joint_is_hard`` for per-joint hard/soft constraint mode (non-cable joints) + - ``vbd:dahl_eps_max`` and ``vbd:dahl_tau`` for optional cable angular Dahl friction Attributes are declared in the ``vbd`` namespace so they can be authored in scenes and in USD as ``newton:vbd:``. + Dahl cable friction is enabled per joint only where both + ``model.vbd.dahl_eps_max`` and ``model.vbd.dahl_tau`` are authored + positive; the attributes default to zero. + Args: builder: Model builder to register attributes on. dahl_defaults_enabled: Deprecated compatibility mode. When True, Dahl parameters - default to positive values. Prefer passing ``False`` and explicitly authoring - positive Dahl values only when Dahl cable friction is desired. + default to positive values instead of zero. + + .. deprecated:: 1.5 + The compatibility mode will be removed; author positive Dahl + values explicitly when Dahl cable friction is desired. """ dahl_eps_default = 0.5 if dahl_defaults_enabled else 0.0 dahl_tau_default = 1.0 if dahl_defaults_enabled else 0.0 if dahl_defaults_enabled: warnings.warn( - "Implicit positive Dahl defaults in SolverVBD.register_custom_attributes() are deprecated " - "and will be disabled by default in a future release. Pass dahl_defaults_enabled=False and " - "explicitly author positive model.vbd.dahl_eps_max and model.vbd.dahl_tau values to enable " + "SolverVBD.register_custom_attributes(dahl_defaults_enabled=True) is deprecated " + "and the compatibility mode will be removed in a future release. Explicitly author " + "positive model.vbd.dahl_eps_max and model.vbd.dahl_tau values to enable " "Dahl cable friction.", DeprecationWarning, stacklevel=2, @@ -1598,8 +1691,8 @@ def set_rigid_history_update(self, update: bool): kinematic flag) while update is disabled: the per-body lists depend on effective inverse mass and are not rebuilt until the next refresh. - Joint AVBD maintenance (C0 snapshot, lambda decay) - runs every step regardless of this flag via step_joint_C0_lambda(). + Joint AVBD maintenance (C0 snapshot, lambda decay, adaptive penalty + upkeep) runs every step regardless of this flag via step_joint_C0_lambda(). Rigid contact history snapshotting also runs every step when enabled. This setting applies only to the next call to :meth:`step` and is then @@ -1612,23 +1705,25 @@ def set_rigid_history_update(self, update: bool): self._update_rigid_history = update def set_joint_constraint_mode(self, joint_index: int, hard: bool, slot: int | None = None): - """Set hard or soft constraint mode for a joint's structural slots at runtime. + """Set hard or soft constraint mode for a joint's structural slots. Hard mode (augmented Lagrangian): uses persistent lambda + C0 stabilization to drive constraint violation toward zero across iterations. Soft mode (penalty-only): uses penalty stiffness only (no lambda or C0 state). - Structural slots are LINEAR (slot 0) and ANGULAR (slot 1). Drive/limit slots - (slot 2+) are always soft and cannot be set to hard. + Non-cable structural slots are LINEAR (slot 0) and ANGULAR (slot 1). + Builder-created cable joints expose STRETCH (slot 0), SHEAR + (slot 1), BEND (slot 2), and TWIST (slot 3). Other drive/limit slots + are always soft and cannot be set to hard. - By default, cable stretch and bend slots are soft, while non-cable - structural slots are hard. + By default, cable stretch, shear, bend, and twist slots are soft, while + non-cable structural slots are hard. - Hard/soft mode can also be authored per joint at build time via the - ``vbd:joint_is_hard`` custom attribute, avoiding a runtime - :meth:`set_joint_constraint_mode` call:: + For non-cable joints, hard/soft mode can also be authored per joint at + build time via the ``vbd:joint_is_hard`` custom attribute, avoiding a + runtime :meth:`set_joint_constraint_mode` call:: - SolverVBD.register_custom_attributes(builder, dahl_defaults_enabled=False) # before adding joints + SolverVBD.register_custom_attributes(builder) # before adding joints builder.add_joint_fixed(..., custom_attributes={"vbd:joint_is_hard": 0}) model = builder.finalize() solver = SolverVBD(model, ...) @@ -1637,13 +1732,13 @@ def set_joint_constraint_mode(self, joint_index: int, hard: bool, slot: int | No joint_index: Index of the joint to modify. hard: True for hard mode (AL), False for soft mode (penalty-only). slot: Specific slot index to set. If None, sets all structural slots. - Use JointSlot.LINEAR / JointSlot.ANGULAR (equivalently - JointSlot.STRETCH / JointSlot.BEND for cables). + Use JointSlot.LINEAR / JointSlot.ANGULAR for non-cable joints, + or JointSlot.STRETCH / JointSlot.SHEAR / JointSlot.BEND / + JointSlot.TWIST for cables. Raises: - ValueError: If the joint index is out of range, or the slot is a - drive/limit slot (>= 2), or the slot exceeds the joint's - constraint dimension. + ValueError: If the joint index is out of range or the slot is not a + structural slot for this joint. """ n_j = self.model.joint_count if joint_index < 0 or joint_index >= n_j: @@ -1653,44 +1748,36 @@ def set_joint_constraint_mode(self, joint_index: int, hard: bool, slot: int | No c_start_np = self._to_numpy(self.joint_constraint_start, dtype=np.int32) c_dim_np = self._to_numpy(self.joint_constraint_dim, dtype=np.int32) is_hard_np = self._to_numpy(self.joint_is_hard, dtype=np.int32) + joint_type_np = self._to_numpy(self.model.joint_type, dtype=np.int32) c0 = int(c_start_np[joint_index]) cdim = int(c_dim_np[joint_index]) + joint_type = int(joint_type_np[joint_index]) + structural_count = cdim if joint_type == int(JointType.CABLE) else min(cdim, 2) val = 1 if hard else 0 if slot is not None: - if slot < 0 or slot >= 2: - raise ValueError( - f"Cannot set hard mode on slot={slot}. " - "Only structural slots (LINEAR=0, ANGULAR=1) support hard mode." - ) - if slot >= cdim: + if slot < 0 or slot >= structural_count: + if structural_count == 0: + names = "no structural slots" + elif joint_type == int(JointType.CABLE): + names = "STRETCH=0, SHEAR=1, BEND=2, TWIST=3" + elif structural_count == 1: + names = "LINEAR=0" + else: + names = "LINEAR=0, ANGULAR=1" raise ValueError( - f"slot={slot} exceeds joint constraint dimension ({cdim}) for joint_index={joint_index}." + f"Cannot set hard mode on slot={slot}: this joint has " + f"{structural_count} structural slot(s) ({names})." ) is_hard_np[c0 + slot] = val else: - structural_count = min(cdim, 2) for s in range(structural_count): is_hard_np[c0 + s] = val - self.joint_is_hard = wp.array(is_hard_np, dtype=wp.int32, device=self.device) - - if not hard: - lam_lin_np = self._to_numpy(self.joint_lambda_lin) - lam_ang_np = self._to_numpy(self.joint_lambda_ang) - C0_lin_np = self._to_numpy(self.joint_C0_lin) - C0_ang_np = self._to_numpy(self.joint_C0_ang) - if slot is None or slot == 0: - lam_lin_np[joint_index] = [0.0, 0.0, 0.0] - C0_lin_np[joint_index] = [0.0, 0.0, 0.0] - if (slot is None or slot == 1) and cdim > 1: - lam_ang_np[joint_index] = [0.0, 0.0, 0.0] - C0_ang_np[joint_index] = [0.0, 0.0, 0.0] - self.joint_lambda_lin = wp.array(lam_lin_np, dtype=wp.vec3, device=self.device) - self.joint_lambda_ang = wp.array(lam_ang_np, dtype=wp.vec3, device=self.device) - self.joint_C0_lin = wp.array(C0_lin_np, dtype=wp.vec3, device=self.device) - self.joint_C0_ang = wp.array(C0_ang_np, dtype=wp.vec3, device=self.device) + # Mutate in place: a rebuilt wp.array would orphan pointers captured + # in existing CUDA graphs, silently ignoring the mode change on replay. + self.joint_is_hard.assign(is_hard_np) @override def step( @@ -1725,7 +1812,7 @@ def step( Raises: RuntimeError: If required rigid contact-matching data is unavailable, or contact-history storage would - need to be allocated or grown during CUDA graph capture. + need to be allocated or grown during graph capture. """ self._apply_module_options() update_rigid = self._update_rigid_history @@ -1743,9 +1830,7 @@ def step( # Snapshot solved rigid contact state for next-frame warm-start. self._snapshot_rigid_contact_history(contacts) - self._finalize_rigid_bodies( - state_in, state_out, dt, apply_stick_deadzone=contacts is not None and self.rigid_contact_hard - ) + self._finalize_rigid_bodies(state_in, state_out, dt) self._finalize_particles(state_out, dt) @override @@ -1785,12 +1870,14 @@ def reset( rigid :meth:`step` consumes the pose and cable rebaseline even when ``contacts=None``, so author the final pose (or run :func:`~newton.eval_fk`) before stepping; contact invalidation instead waits for a fresh refresh. - With ``rigid_contact_history=True``, only ``contact_matching="latest"`` - is supported; VBD cannot invalidate sticky matcher state owned by the - collision pipeline. Reset does not change ``set_rigid_history_update()``; - leave rigid-history refresh enabled for the next contact-bearing step. - Reusing contacts (``set_rigid_history_update(False)``) is unsupported - only while contact invalidation is still pending. + VBD cold-starts its numeric contact state for reset-selected worlds. + Frame-to-frame correspondence and sticky contact geometry remain owned by + :class:`~newton.CollisionPipeline`; construct a new pipeline to discard + that history after a discontinuous episode reset. Reset does not change + ``set_rigid_history_update()``; leave rigid history refresh enabled for the + next contact-bearing step. Reusing contacts + (``set_rigid_history_update(False)``) is unsupported only while contact + invalidation is still pending. Args: state: The simulation state to reset (modified in place). @@ -1910,23 +1997,13 @@ def _snapshot_rigid_contact_history(self, contacts: Contacts | None): dim=contact_launch_dim, inputs=[ contacts.rigid_contact_count, - contacts.rigid_contact_point0, - contacts.rigid_contact_point1, - contacts.rigid_contact_offset0, - contacts.rigid_contact_offset1, contacts.rigid_contact_normal, self.body_body_contact_lambda, - self.body_body_contact_stick_flag, self.body_body_contact_penalty_k, ], outputs=[ self._prev_contact_lambda, - self._prev_contact_stick_flag, self._prev_contact_penalty_k, - self._prev_contact_point0, - self._prev_contact_point1, - self._prev_contact_offset0, - self._prev_contact_offset1, self._prev_contact_normal, ], device=self.device, @@ -2057,11 +2134,17 @@ def _initialize_rigid_bodies( internal_rigid = model.body_count > 0 and not self.integrate_with_external_rigid_solver rigid_capacity = contacts.rigid_contact_max if contacts is not None else 0 + # Rigid contact history is cross-replay-persistent state: allocating it + # during capture records a `wp.zeros` fill into the graph, which then + # re-zeros the warm-start buffers on every replay -- silently + # equivalent to `rigid_contact_history=False`. So this guard fires + # unconditionally when capturing, regardless of the device's + # allocation-during-capture support. if self.device.is_capturing and internal_rigid and self.rigid_contact_history: history_capacity = 0 if self._prev_contact_lambda is None else self._prev_contact_lambda.shape[0] if history_capacity < rigid_capacity: raise RuntimeError( - "SolverVBD contact history must be allocated before CUDA graph capture. " + "SolverVBD contact history must be allocated before graph capture. " "Construct CollisionPipeline before SolverVBD, or run one uncaptured solver step before capture." ) @@ -2127,9 +2210,9 @@ def _initialize_rigid_bodies( if contacts.rigid_contact_match_index is None: raise RuntimeError( "SolverVBD(rigid_contact_history=True) requires Contacts with " - "rigid_contact_match_index populated. Create contacts through " - 'CollisionPipeline(contact_matching="latest") for VBD warm-starting, ' - "or set rigid_contact_history=False." + "rigid_contact_match_index populated. Use " + 'CollisionPipeline(contact_matching="latest") or ' + 'CollisionPipeline(contact_matching="sticky"), or set rigid_contact_history=False.' ) history_required = contact_launch_dim @@ -2140,12 +2223,7 @@ def _initialize_rigid_bodies( history = RigidContactHistory() history.lambda_ = self._prev_contact_lambda - history.stick_flag = self._prev_contact_stick_flag history.penalty_k = self._prev_contact_penalty_k - history.point0 = self._prev_contact_point0 - history.point1 = self._prev_contact_point1 - history.offset0 = self._prev_contact_offset0 - history.offset1 = self._prev_contact_offset1 history.normal = self._prev_contact_normal wp.launch( @@ -2170,10 +2248,6 @@ def _initialize_rigid_bodies( self.rigid_contact_k_start_value, ], outputs=[ - contacts.rigid_contact_point0, - contacts.rigid_contact_point1, - contacts.rigid_contact_offset0, - contacts.rigid_contact_offset1, self.body_body_contact_penalty_k, self.body_body_contact_lambda, self.body_body_contact_material_kd, @@ -2248,7 +2322,6 @@ def _initialize_rigid_bodies( ], device=self.device, ) - self.body_body_contact_stick_flag.zero_() # Accumulate joint_f into body wrenches (scratch buffer avoids mutating user state). body_f_for_integration = state_in.body_f @@ -2311,11 +2384,14 @@ def _initialize_rigid_bodies( kernel=step_joint_C0_lambda, dim=model.joint_count, inputs=[ + model.joint_type, model.joint_enabled, model.joint_parent, model.joint_child, model.joint_X_p, model.joint_X_c, + self.joint_cable_rest_kb_local, + self.joint_cable_rest_twist, self.body_q_prev, model.body_q, self.joint_constraint_start, @@ -2336,7 +2412,7 @@ def _initialize_rigid_bodies( device=self.device, ) - # Compute Dahl hysteresis parameters for cable bending (once per timestep, frozen during iterations) + # Compute cable bend/twist Dahl hysteresis parameters once per timestep. if self.enable_dahl_friction and model.joint_count > 0: wp.launch( kernel=compute_cable_dahl_parameters, @@ -2351,8 +2427,10 @@ def _initialize_rigid_bodies( model.joint_X_c, self.joint_constraint_start, self.joint_penalty_k_max, + self.joint_is_hard, + self.joint_cable_rest_kb_local, + self.joint_cable_rest_twist, self.body_q_prev, - model.body_q, self.joint_sigma_prev, self.joint_kappa_prev, self.joint_dkappa_prev, @@ -2635,7 +2713,12 @@ def _solve_particle_iteration( wp.copy(state_out.particle_q, state_in.particle_q) def _solve_rigid_body_iteration( - self, state_in: State, state_out: State, control: Control, contacts: Contacts | None, dt: float + self, + state_in: State, + state_out: State, + control: Control, + contacts: Contacts | None, + dt: float, ): """Solve one AVBD iteration for rigid bodies (per-iteration phase). @@ -2800,6 +2883,8 @@ def _solve_rigid_body_iteration( model.joint_X_p, model.joint_X_c, model.joint_axis, + self.joint_cable_rest_kb_local, + self.joint_cable_rest_twist, model.joint_qd_start, model.joint_target_q_start, self.joint_constraint_start, @@ -2858,17 +2943,12 @@ def _solve_rigid_body_iteration( self.body_body_contact_material_mu, self.body_body_contact_C0, self.rigid_contact_alpha, - self.rigid_contact_stick_motion_eps, self.rigid_contact_hard, - self.body_inv_mass_effective, self.body_body_contact_material_ke, self.rigid_linear_beta, self.body_body_contact_penalty_k, # input/output self.body_body_contact_lambda, # input/output ], - outputs=[ - self.body_body_contact_stick_flag, - ], device=self.device, ) @@ -2908,6 +2988,8 @@ def _solve_rigid_body_iteration( model.joint_X_p, model.joint_X_c, model.joint_axis, + self.joint_cable_rest_kb_local, + self.joint_cable_rest_twist, model.joint_qd_start, model.joint_target_q_start, self.joint_constraint_start, @@ -3095,13 +3177,11 @@ def _finalize_particles(self, state_out: State, dt: float): device=self.device, ) - def _finalize_rigid_bodies(self, state_in: State, state_out: State, dt: float, apply_stick_deadzone: bool): + def _finalize_rigid_bodies(self, state_in: State, state_out: State, dt: float): """Finalize rigid body velocities and Dahl friction state after AVBD iterations (post-iteration phase). - Updates rigid body velocities using BDF1 and updates Dahl hysteresis state for cable bending. - Also transfers the final body poses from state_in to state_out. When requested, - the fused finalize kernel first applies the body-level stick-contact deadzone - before computing velocity from the accepted pose. + Updates rigid body velocities using BDF1 and updates Dahl hysteresis state for cable bend/twist. + Also transfers the final body poses from state_in to state_out. """ model = self.model @@ -3115,13 +3195,6 @@ def _finalize_rigid_bodies(self, state_in: State, state_out: State, dt: float, a dt, state_in.body_q, model.body_com, - self.body_body_contact_buffer_pre_alloc, - self.body_body_contact_counts, - self.body_body_contact_indices, - self.body_body_contact_stick_flag, - int(apply_stick_deadzone), - self.rigid_contact_stick_freeze_translation_eps, - self.rigid_contact_stick_freeze_angular_eps, ], outputs=[self.body_q_prev, state_out.body_qd, state_in.body_qd, state_out.body_q], dim=model.body_count, @@ -3141,8 +3214,9 @@ def _finalize_rigid_bodies(self, state_in: State, state_out: State, dt: float, a self.joint_constraint_start, self.joint_penalty_k_max, self.joint_is_hard, + self.joint_cable_rest_kb_local, + self.joint_cable_rest_twist, state_out.body_q, - model.body_q, self.joint_dahl_eps_max, self.joint_dahl_tau, self.joint_sigma_prev, diff --git a/newton/_src/usd/utils.py b/newton/_src/usd/utils.py index d8e43023bb..dab81d6bbe 100644 --- a/newton/_src/usd/utils.py +++ b/newton/_src/usd/utils.py @@ -857,6 +857,143 @@ def _expand_indexed_primvar( return values[indices] +def _split_corners_into_vertices( + points: np.ndarray, + indices: np.ndarray, + corner_dirs: np.ndarray, + corner_uvs: np.ndarray | None, + angle_threshold_deg: float, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]: + """Duplicate vertices whose faceVarying corners disagree in normal direction or UV. + + The corners of a vertex are clustered greedily in corner order: a corner joins the + first cluster whose mean direction is within ``angle_threshold_deg`` and whose UV + matches, otherwise it starts a new cluster. Every cluster becomes one output vertex, + numbered by the corner that created it. + + Most vertices resolve to a single cluster, which is decided for all of them at once + with array operations; only the vertices that fail that test run the sequential + clustering. A vertex whose corner directions all lie within half the threshold angle + of their mean lie within the full threshold of each other, so the sequential pass + would place every one of them in the first cluster. + + A corner whose angle to the cluster mean is exactly ``angle_threshold_deg`` may fall + on either side of the comparison: the dot product is rounded differently depending on + the arithmetic used, so which cluster such a corner lands in is not defined beyond + "one of the clusters it is within the threshold of". + + Args: + points: Source vertex positions, shape [vertex_count, 3]. + indices: Face-corner vertex indices, shape [corner_count]. + corner_dirs: Unit corner normals, shape [corner_count, 3]. + corner_uvs: Per-corner UVs, shape [corner_count, channels], or ``None``. + angle_threshold_deg: Maximum angle between a corner normal and its cluster mean [deg]. + + Returns: + Split positions, the remapped corner indices, the per-vertex normals, and the + per-vertex UVs (``None`` when ``corner_uvs`` is ``None``). + """ + corner_count = len(indices) + if corner_count == 0: + empty_uvs = None if corner_uvs is None else corner_uvs[:0] + return points[:0], indices.copy(), np.zeros((0, 3), dtype=np.float32), empty_uvs + + cos_thresh = math.cos(math.radians(angle_threshold_deg)) + cos_half_thresh = math.cos(math.radians(angle_threshold_deg) * 0.5) + + # Group the corners by vertex; the stable sort keeps corner order inside each group. + order = np.argsort(indices, kind="stable") + grouped_vertices = indices[order] + grouped_dirs = corner_dirs[order] + grouped_uvs = None if corner_uvs is None else corner_uvs[order] + + starts = np.flatnonzero(np.concatenate(([True], grouped_vertices[1:] != grouped_vertices[:-1]))) + group_sizes = np.diff(np.append(starts, corner_count)) + group_vertices = grouped_vertices[starts] + group_of_corner = np.repeat(np.arange(len(starts)), group_sizes) + + # The normalized sum of a group's corner directions is the cluster mean it would end + # up with if every corner joined the same cluster. + dir_sums = np.add.reduceat(grouped_dirs, starts, axis=0) + dir_means = dir_sums / np.clip(np.linalg.norm(dir_sums, axis=1, keepdims=True), 1e-30, None) + dots = np.einsum("ij,ij->i", grouped_dirs, dir_means[group_of_corner]) + single_cluster = np.minimum.reduceat(dots, starts) >= cos_half_thresh + if grouped_uvs is not None: + same_uv = np.all(grouped_uvs == grouped_uvs[starts][group_of_corner], axis=1) + single_cluster &= np.logical_and.reduceat(same_uv, starts) + + # Provisional ids: the single-cluster groups first, then the sequential pass. + simple_groups = np.flatnonzero(single_cluster) + provisional_of_group = np.full(len(starts), -1, dtype=np.int64) + provisional_of_group[simple_groups] = np.arange(len(simple_groups)) + provisional_of_corner = provisional_of_group[group_of_corner] + + next_id = len(simple_groups) + split_creation: list[int] = [] + split_vertices: list[int] = [] + split_dir_sums: list[tuple[float, float, float]] = [] + split_uvs: list[tuple[float, ...]] = [] + for group in np.flatnonzero(~single_cluster): + begin = int(starts[group]) + end = begin + int(group_sizes[group]) + source_vertex = int(group_vertices[group]) + clusters: list[list] = [] + for corner in range(begin, end): + dir_x, dir_y, dir_z = (float(value) for value in grouped_dirs[corner]) + corner_uv = None if grouped_uvs is None else tuple(grouped_uvs[corner].tolist()) + for cluster in clusters: + sum_x, sum_y, sum_z = cluster[0], cluster[1], cluster[2] + # Scalar arithmetic rounds this dot product differently from the + # equivalent NumPy expression (which uses FMA and a scaled norm), so a + # corner sitting exactly on the threshold can cluster either way. + scale = max(math.sqrt(sum_x * sum_x + sum_y * sum_y + sum_z * sum_z), 1e-30) + if (sum_x * dir_x + sum_y * dir_y + sum_z * dir_z) / scale < cos_thresh: + continue + if corner_uv is not None and cluster[3] != corner_uv: + continue + cluster[0] = sum_x + dir_x + cluster[1] = sum_y + dir_y + cluster[2] = sum_z + dir_z + provisional_of_corner[corner] = cluster[4] + break + else: + clusters.append([dir_x, dir_y, dir_z, corner_uv, next_id]) + provisional_of_corner[corner] = next_id + split_creation.append(int(order[corner])) + split_vertices.append(source_vertex) + if corner_uv is not None: + split_uvs.append(corner_uv) + next_id += 1 + # Cluster order matches the id order assigned above. + split_dir_sums.extend((cluster[0], cluster[1], cluster[2]) for cluster in clusters) + + creation_corners = np.concatenate([order[starts[simple_groups]], np.asarray(split_creation, dtype=np.int64)]) + source_vertices = np.concatenate([group_vertices[simple_groups], np.asarray(split_vertices, dtype=np.int64)]) + new_dir_sums = np.concatenate( + [dir_sums[simple_groups], np.asarray(split_dir_sums, dtype=np.float64).reshape(-1, 3)] + ) + + # Number the output vertices by the corner that created them, matching the order a + # purely sequential pass over the corners would produce. + final_of_provisional = np.empty(len(creation_corners), dtype=np.int64) + final_of_provisional[np.argsort(creation_corners, kind="stable")] = np.arange(len(creation_corners)) + provisional_in_final_order = np.argsort(final_of_provisional, kind="stable") + + new_indices = np.empty(corner_count, dtype=indices.dtype) + new_indices[order] = final_of_provisional[provisional_of_corner] + new_points = points[source_vertices[provisional_in_final_order]] + ordered_sums = new_dir_sums[provisional_in_final_order] + lengths = np.clip(np.linalg.norm(ordered_sums, axis=1, keepdims=True), 1e-30, None) + new_normals = (ordered_sums / lengths).astype(np.float32) + + new_uvs = None + if grouped_uvs is not None: + split_uv_rows = np.asarray(split_uvs, dtype=grouped_uvs.dtype).reshape(-1, grouped_uvs.shape[1]) + new_uvs = np.concatenate([grouped_uvs[starts[simple_groups]], split_uv_rows])[provisional_in_final_order] + + return new_points, new_indices, new_normals, new_uvs + + def _triangulate_face_varying_indices(counts: Sequence[int], flip_winding: bool) -> np.ndarray: """Return flattened corner indices for fan-triangulated face-varying data.""" counts_i32 = np.asarray(counts, dtype=np.int32) @@ -1065,6 +1202,121 @@ def _get_mesh_from_source( ) +def _material_surface_shader(material: UsdShade.Material | None) -> UsdShade.Shader | None: + """Return the surface shader driving a material (UsdPreviewSurface or MDL).""" + if not material: + return None + surface_output = material.GetSurfaceOutput() or material.GetOutput("surface") or material.GetOutput("mdl:surface") + if surface_output: + source = surface_output.GetConnectedSource() + if source: + return UsdShade.Shader(source[0].GetPrim()) + for child in material.GetPrim().GetChildren(): + if child.IsA(UsdShade.Shader): + return UsdShade.Shader(child) + return None + + +def _uvtexture_reader_varname(texture_shader: UsdShade.Shader) -> str | None: + """Return the primvar name a ``UsdUVTexture`` reads via its ``st`` -> ``UsdPrimvarReader_float2``.""" + st_input = texture_shader.GetInput("st") + source = st_input.GetConnectedSource() if st_input else None + if not source: + return None + reader = UsdShade.Shader(source[0].GetPrim()) + varname = reader.GetInput("varname") + if not varname: + return None + value = varname.Get() + if value is None: + try: + attrs = UsdShade.Utils.GetValueProducingAttributes(varname) + except Exception: + attrs = () + value = attrs[0].Get() if attrs else None + return str(value) if value else None + + +def _uv_primvar_name_from_shader(shader: UsdShade.Shader | None) -> str | None: + """Resolve the texcoord primvar name a shader's base-color texture reads, or ``None``. + + - ``UsdPreviewSurface``: follow the base-color ``UsdUVTexture``'s ``st`` input to its + ``UsdPrimvarReader_float2`` and read ``inputs:varname``. + - ``OmniPBR`` and other MDL shaders: ``st_``. + """ + if shader is None: + return None + try: + shader_id = shader.GetIdAttr().Get() + except Exception: + shader_id = None + if shader_id == "UsdPreviewSurface": + for color_name in ("baseColor", "diffuseColor"): + color_input = shader.GetInput(color_name) + source = color_input.GetConnectedSource() if color_input else None + if not source: + continue + texture = UsdShade.Shader(source[0].GetPrim()) + if texture.GetIdAttr().Get() == "UsdUVTexture": + name = _uvtexture_reader_varname(texture) + if name: + return name + uv_index_input = shader.GetInput("uv_space_index") + if uv_index_input: + value = uv_index_input.Get() + if value is not None: + return f"st_{int(value)}" + return None + + +def _resolve_material_uv_primvar_name(prim: Usd.Prim) -> str | None: + """Resolve the texcoord primvar name from the material(s) bound to a mesh or its subsets.""" + candidate_prims = [prim] + try: + subsets = UsdShade.MaterialBindingAPI(prim).GetMaterialBindSubsets() + candidate_prims += [subset.GetPrim() for subset in subsets] + except Exception: + pass + for candidate in candidate_prims: + name = _uv_primvar_name_from_shader(_material_surface_shader(_get_bound_material(candidate))) + if name: + return name + return None + + +def _find_uv_primvar(prim: Usd.Prim): + """Return a mesh's texture-coordinate primvar, or ``None``. + + The primvar name is resolved from the bound material's shader network — the + ``UsdPreviewSurface`` texture reader's ``varname`` or an MDL/OmniPBR + ``uv_space_index`` — because a mesh may carry several UV sets and only the + material identifies the correct one. Falls back to the conventional ``st`` + primvar, then to the first ``float2``/``texCoord2f`` primvar named ``st*``/``uv*``. + """ + api = UsdGeom.PrimvarsAPI(prim) + + resolved = _resolve_material_uv_primvar_name(prim) + if resolved: + primvar = api.GetPrimvar(resolved) + if primvar and primvar.HasValue(): + return primvar + + primvar = api.GetPrimvar("st") + if primvar and primvar.HasValue(): + return primvar + + fallback = None + for candidate in api.GetPrimvarsWithValues(): + if candidate.GetTypeName() not in (Sdf.ValueTypeNames.TexCoord2fArray, Sdf.ValueTypeNames.Float2Array): + continue + name = candidate.GetPrimvarName().lower() + if name.startswith("st"): + return candidate + if (name.startswith("uv") or name == "map1") and fallback is None: + fallback = candidate + return fallback + + @overload def get_mesh( source: Usd.Prim | Usd.Stage | str | os.PathLike[str], @@ -1288,7 +1540,7 @@ def get_mesh( # faceVarying normal conversion, so we don't split again in the UV pass. did_split_vertices = False if load_uvs: - uv_primvar = UsdGeom.PrimvarsAPI(prim).GetPrimvar("st") + uv_primvar = _find_uv_primvar(prim) if uv_primvar: uvs = uv_primvar.Get() if uvs is not None: @@ -1358,68 +1610,33 @@ def get_mesh( nlen = np.clip(nlen, 1e-30, None) Ndir = Nfv / nlen - cos_thresh = np.cos(np.deg2rad(vertex_splitting_angle_threshold_deg)) - - # For each original vertex v, we'll keep a list of clusters: - # each cluster stores (sum_dir, count, new_vid) - clusters_per_v = [[] for _ in range(V)] - - new_points = [] - new_norm_sums = [] # accumulate directions per new vertex id - new_indices = np.empty_like(indices) - new_uvs = [] if uvs is not None else None - - # Helper to create a new vertex clone from original v - def _new_vertex_from(v, n_dir, corner_idx): - new_vid = len(new_points) - new_points.append(points[v]) - new_norm_sums.append(n_dir.copy()) - clusters_per_v[v].append([n_dir.copy(), 1, new_vid]) - if new_uvs is not None: - # Use corner UV if faceVarying, otherwise use vertex UV - if uvs_interpolation == UsdGeom.Tokens.faceVarying: - new_uvs.append(uvs[corner_idx]) - else: - new_uvs.append(uvs[v]) - return new_vid - - # Assign each corner to a cluster (new vertex) based on angular proximity - for c in range(C): - v = int(indices[c]) - n_dir = Ndir[c] - - clusters = clusters_per_v[v] - assigned = False - # try to match an existing cluster - for cl in clusters: - sum_dir, cnt, new_vid = cl - # compare with current mean direction (sum_dir normalized) - mean_dir = sum_dir / max(np.linalg.norm(sum_dir), 1e-30) - if float(np.dot(mean_dir, n_dir)) >= cos_thresh: - # assign to this cluster - cl[0] = sum_dir + n_dir - cl[1] = cnt + 1 - new_norm_sums[new_vid] += n_dir - new_indices[c] = new_vid - assigned = True - break - - if not assigned: - new_vid = _new_vertex_from(v, n_dir, c) - new_indices[c] = new_vid - - new_points = np.asarray(new_points, dtype=np.float64) - - # Produce per-vertex normalized normals for the new vertices - new_norm_sums = np.asarray(new_norm_sums, dtype=np.float64) - nn = np.linalg.norm(new_norm_sums, axis=1, keepdims=True) - nn = np.clip(nn, 1e-30, None) - new_vertex_normals = (new_norm_sums / nn).astype(np.float32) - - points = new_points - indices = new_indices - normals = new_vertex_normals - uvs = new_uvs + # faceVarying UVs carry one value per corner; if the count does + # not match, they can't be indexed per-corner, so drop them + # (matching the non-splitting UV path below). + uvs_facevarying = uvs is not None and uvs_interpolation == UsdGeom.Tokens.faceVarying + if uvs_facevarying and len(uvs) != C: + logger.info( + "Mesh %s: UV primvar length (%d) does not match corner count (%d); dropping UVs.", + prim.GetPath(), + len(uvs), + C, + ) + uvs = None + uvs_facevarying = False + + # Corners that share a smooth normal but carry different faceVarying UVs + # lie on a texture seam and must not be merged, or the seam's UVs would + # collapse onto one value. + if uvs is None: + corner_uvs = None + else: + corner_uvs = np.asarray(uvs).reshape(len(uvs), -1) + if not uvs_facevarying: + corner_uvs = corner_uvs[indices] + + points, indices, normals, uvs = _split_corners_into_vertices( + points, indices, Ndir, corner_uvs, vertex_splitting_angle_threshold_deg + ) # Vertex splitting creates a new per-vertex layout (and UVs # if available). Skip the later faceVarying UV split to avoid # dropping/duplicating UVs. @@ -2449,6 +2666,72 @@ def _extract_preview_surface_properties(shader: UsdShade.Shader | None, prim: Us return properties +def _output_channel_count(type_name: Sdf.ValueTypeName) -> int: + """Return the component count of a shader output value type (float -> 1, float3 -> 3).""" + default = type_name.defaultValue + if hasattr(default, "__len__"): + return len(default) + return 1 + + +# Base-color input name fragments used to identify the diffuse/albedo parameter +# feeding a surface shader. A texture connection's shape alone is ambiguous +# (a normal map is a 3-channel ``rgb`` connection like a diffuse map), so the +# input name is the disambiguator for both connected and direct-asset textures. +_COLOR_TEXTURE_INPUT_NAMES = ("diffuse", "albedo", "basecolor", "base_color", "displaycolor") + + +def _is_color_texture_input_name(base_name: str) -> bool: + """Return whether an input name denotes a base-color/albedo texture parameter.""" + name = base_name.lower() + return any(fragment in name for fragment in _COLOR_TEXTURE_INPUT_NAMES) + + +def _color_texture_from_input(surface_input: UsdShade.Input, prim: Usd.Prim) -> str | np.ndarray | None: + """Return the base-color texture feeding a surface shader input, if any. + + The input must be a base-color/albedo parameter by name: a normal map is + conventionally wired as ``UsdUVTexture.outputs:rgb`` too, so its 3-channel + connection is indistinguishable from a diffuse map by shape alone. Two + shapes are then supported: + + - A connected ``UsdUVTexture``: accepted only when the connected output is a + multi-channel color output (``rgb`` / ``rgba``). Single-channel outputs + (``r``/``g``/``b``/``a``) are scalar data maps (roughness, metallic, ...) + and are ignored. + - A direct asset value (e.g. an MDL ``diffuse_texture`` parameter): there is + no texture node to inspect, so the base-color input is identified by name. + """ + if not _is_color_texture_input_name(surface_input.GetBaseName()): + return None + try: + producing = UsdShade.Utils.GetValueProducingAttributes(surface_input) + except Exception: + producing = () + for attr in producing: + source_shader = UsdShade.Shader(attr.GetPrim()) + try: + if source_shader.GetIdAttr().Get() != "UsdUVTexture": + continue + except Exception: + continue + if _output_channel_count(attr.GetTypeName()) < 3: + continue + texture = _find_texture_in_shader(source_shader, prim) + if texture is not None: + return texture + + try: + connected = surface_input.HasConnectedSource() + except Exception: + connected = False + if not connected: + asset = surface_input.Get() + if asset: + return _resolve_color_texture_asset(asset, prim, surface_input.GetAttr()) + return None + + def _extract_shader_properties(shader: UsdShade.Shader | None, prim: Usd.Prim) -> dict[str, Any]: """Extract common material properties from a shader node. @@ -2495,19 +2778,10 @@ def _extract_shader_properties(shader: UsdShade.Shader | None, prim: Usd.Prim) - if properties["texture"] is None: for inp in shader.GetInputs(): - name = inp.GetBaseName() - if inp.HasConnectedSource(): - source = inp.GetConnectedSource() - source_shader = UsdShade.Shader(source[0].GetPrim()) - texture = _find_texture_in_shader(source_shader, prim) - if texture is not None: - properties["texture"] = texture - break - elif "file" in name or "texture" in name: - asset = inp.Get() - if asset: - properties["texture"] = _resolve_color_texture_asset(asset, prim, inp.GetAttr()) - break + texture = _color_texture_from_input(inp, prim) + if texture is not None: + properties["texture"] = texture + break return properties diff --git a/newton/_src/utils/benchmark.py b/newton/_src/utils/benchmark.py index 9fe250bf5d..72c6138edb 100644 --- a/newton/_src/utils/benchmark.py +++ b/newton/_src/utils/benchmark.py @@ -175,13 +175,20 @@ def run_benchmark(benchmark_cls, number=1, print_results=True): else: combinations = [()] + cache = None + has_cache = hasattr(benchmark_cls, "setup_cache") + if has_cache: + cache = benchmark_cls().setup_cache() + has_cache = cache is not None + results = {} # For each parameter combination: for params in combinations: + call_params = (cache, *params) if has_cache else params # Create a fresh benchmark instance. instance = benchmark_cls() if hasattr(instance, "setup"): - instance.setup(*params) + instance.setup(*call_params) # Iterate over all attributes to find benchmark methods. for attr in dir(instance): if attr.startswith("time_") or attr.startswith("track_"): @@ -190,24 +197,24 @@ def run_benchmark(benchmark_cls, number=1, print_results=True): samples = [] if attr.startswith("time_"): # Warmup run (not measured). - method(*params) + method(*call_params) wp.synchronize() # Run timing benchmarks multiple times and measure elapsed time. for _ in range(number): start = time.perf_counter() - method(*params) + method(*call_params) t = time.perf_counter() - start samples.append(t) elif attr.startswith("track_"): # Run tracking benchmarks multiple times and record returned values. for _ in range(number): - val = method(*params) + val = method(*call_params) samples.append(val) # Compute the average result. avg = sum(samples) / len(samples) results[(attr, params)] = avg if hasattr(instance, "teardown"): - instance.teardown(*params) + instance.teardown(*call_params) if print_results: print("\n=== Benchmark Results ===") diff --git a/newton/_src/utils/cable.py b/newton/_src/utils/cable.py index 5c0ae322d2..d4894d8a5f 100644 --- a/newton/_src/utils/cable.py +++ b/newton/_src/utils/cable.py @@ -5,38 +5,139 @@ import math from collections.abc import Sequence +from typing import NamedTuple, overload import warp as wp from ..math import quat_between_vectors_robust +class CableStiffness(NamedTuple): + """Per-joint Kirchhoff rod stiffness for a circular isotropic cross-section. + + Returned by :func:`create_cable_stiffness_from_elastic_moduli` when the + caller supplies either ``poissons_ratio`` or ``shear_modulus``. + + Fields: + + * ``stretch`` -- axial stiffness ``E * A / L`` [N/m] + * ``bend`` -- bending stiffness ``E * I / L`` [N*m / rad] + * ``twist`` -- torsional stiffness ``G * J / L`` [N*m / rad] + + For a circular cross-section the two bending axes are equivalent + (``EI1 == EI2 == EI``); the single ``bend`` field is used for both axes + when assembling the per-joint cable stiffness vector. + + No ``shear`` field, by design. Set a sufficiently large finite + ``shear_stiffness`` separately to approximate an unshearable Kirchhoff rod. + + Being a :class:`typing.NamedTuple`, instances support both attribute + access (``stiffness.bend``) and tuple unpacking + (``stretch, bend, twist = stiffness``). + """ + + stretch: float + bend: float + twist: float + + +@overload +def create_cable_stiffness_from_elastic_moduli( + youngs_modulus: float, + radius: float, + segment_length: float, +) -> tuple[float, float]: ... + + +@overload def create_cable_stiffness_from_elastic_moduli( youngs_modulus: float, radius: float, segment_length: float, -) -> tuple[float, float]: - """Create per-joint rod/cable stiffness parameters from elastic moduli. + *, + poissons_ratio: float, +) -> CableStiffness: ... - For a circular cross-section, this computes material stiffnesses and converts them to the - per-joint stiffness values expected by ``ModelBuilder.add_rod()`` and - ``ModelBuilder.add_rod_graph()``: - - stretch_stiffness = E * A / L [N/m] - - bend_stiffness = E * I / L [N*m] +@overload +def create_cable_stiffness_from_elastic_moduli( + youngs_modulus: float, + radius: float, + segment_length: float, + *, + shear_modulus: float, +) -> CableStiffness: ... - where: - - A = pi * r^2 - - I = (pi * r^4) / 4 (area moment of inertia for a solid circular rod) - - L = segment_length + +def create_cable_stiffness_from_elastic_moduli( + youngs_modulus: float, + radius: float, + segment_length: float, + *, + poissons_ratio: float | None = None, + shear_modulus: float | None = None, +) -> tuple[float, float] | CableStiffness: + """Create per-joint rod/cable stiffness from elastic moduli. + + For a circular cross-section, this computes material stiffnesses and + converts them to the per-joint stiffness values expected by + :meth:`ModelBuilder.add_rod` and :meth:`ModelBuilder.add_rod_graph`: + + * ``stretch = E * A / L`` [N/m] + * ``bend = E * I / L`` [N*m / rad] + * ``twist = G * J / L`` [N*m / rad] (returned only when + ``poissons_ratio`` or ``shear_modulus`` is supplied) + + where ``A = pi * r^2``, ``I = pi * r^4 / 4`` (area moment of inertia about + a diameter), ``J = pi * r^4 / 2`` (polar moment of area), and + ``L = segment_length``. For an isotropic material with Poisson's ratio + ``nu``, the shear modulus is ``G = E / (2 * (1 + nu))``. + + No separate transverse shear stiffness is returned. The split cable API + defaults ``shear_stiffness`` to ``stretch_stiffness`` when omitted; pass an + explicit ``shear_stiffness`` if that default is not desired. The + ``shear_modulus`` keyword supplies ``G`` for torsion/twist only. + + The return shape mirrors what the caller asks for: + + * ``create_cable_stiffness_from_elastic_moduli(E, r, L)`` returns the + plain 2-tuple ``(stretch, bend)`` -- twist is not derivable from + ``E`` alone, so it is omitted. Suitable for stretch-only or + bend-only rods, or when the caller manages ``twist_stiffness`` + separately. + * Supplying ``poissons_ratio`` or ``shear_modulus`` switches the + return to :class:`CableStiffness` with the additional ``twist`` + term. The result both unpacks as a 3-tuple and exposes named + fields ``.stretch``, ``.bend``, ``.twist``. + + When the 2-tuple is passed through the builder without an explicit + ``twist_stiffness``, twist defaults to ``bend`` (the combined-stiffness model). + For material-consistent torsion ``twist / bend = G * J / (E * I) = 1 / (1 + nu)``, + pass ``poissons_ratio`` or ``shear_modulus`` to get the third term. Args: - youngs_modulus: Young's modulus E in Pascals [N/m^2]. - radius: Rod/cable radius r in meters. - segment_length: Segment length L in meters. + youngs_modulus: Young's modulus ``E`` [Pa = N/m^2]. Must be finite + and ``>= 0``. + radius: Rod/cable radius ``r`` [m]. Must be finite and ``> 0``. + segment_length: Per-joint rest length ``L`` [m]. Must be finite and + ``> 0``. + poissons_ratio: Poisson's ratio ``nu`` used to compute the shear + modulus ``G = E / (2 * (1 + nu))``. Keyword-only. Must satisfy + ``-1 < nu < 0.5`` for a stable isotropic 3D material. Mutually + exclusive with ``shear_modulus``. + shear_modulus: Shear modulus ``G`` [Pa]. Keyword-only. Mutually + exclusive with ``poissons_ratio``. Returns: - Tuple `(stretch_stiffness, bend_stiffness)` = `(E*A/L, E*I/L)`. + 2-tuple ``(stretch, bend)`` when neither ``poissons_ratio`` nor + ``shear_modulus`` is supplied; otherwise a :class:`CableStiffness` + NamedTuple ``(stretch, bend, twist)``. + + Raises: + ValueError: if any of ``youngs_modulus``, ``radius``, + ``segment_length``, ``poissons_ratio``, or ``shear_modulus`` is + non-finite or out of range, or if both ``poissons_ratio`` and + ``shear_modulus`` are supplied. """ # Accept ints / numpy scalars, but return plain Python floats. E = float(youngs_modulus) @@ -56,11 +157,37 @@ def create_cable_stiffness_from_elastic_moduli( raise ValueError("radius must be > 0") if L <= 0.0: raise ValueError("segment_length must be > 0") + if poissons_ratio is not None and shear_modulus is not None: + raise ValueError("poissons_ratio and shear_modulus are mutually exclusive") area = math.pi * r * r inertia = 0.25 * math.pi * r**4 - - return E * area / L, E * inertia / L + stretch_stiffness = E * area / L + bend_stiffness = E * inertia / L + + if poissons_ratio is None and shear_modulus is None: + return stretch_stiffness, bend_stiffness + + if shear_modulus is None: + nu = float(poissons_ratio) + if not math.isfinite(nu): + raise ValueError("poissons_ratio must be finite") + if nu <= -1.0 or nu >= 0.5: + raise ValueError("poissons_ratio must satisfy -1 < nu < 0.5") + G = E / (2.0 * (1.0 + nu)) + else: + G = float(shear_modulus) + if not math.isfinite(G): + raise ValueError("shear_modulus must be finite") + if G < 0.0: + raise ValueError("shear_modulus must be >= 0") + + polar_inertia = 0.5 * math.pi * r**4 + return CableStiffness( + stretch=stretch_stiffness, + bend=bend_stiffness, + twist=G * polar_inertia / L, + ) def create_straight_cable_points( diff --git a/newton/_src/utils/download_assets.py b/newton/_src/utils/download_assets.py index bde9d495e2..7e1e86830f 100644 --- a/newton/_src/utils/download_assets.py +++ b/newton/_src/utils/download_assets.py @@ -23,7 +23,7 @@ NEWTON_ASSETS_REF = "261cd1f429619d8ef4f546bd788ab9dea906b5e1" MENAGERIE_URL = "https://github.com/google-deepmind/mujoco_menagerie.git" -MENAGERIE_REF = "feadf76d42f8a2162426f7d226a3b539556b3bf5" +MENAGERIE_REF = "affef0836947b64cc06c4ab1cbf0152835693374" _SHA_RE = re.compile(r"[0-9a-f]{40}") diff --git a/newton/_src/utils/heightfield.py b/newton/_src/utils/heightfield.py index 2a24650dc0..f04e2a7469 100644 --- a/newton/_src/utils/heightfield.py +++ b/newton/_src/utils/heightfield.py @@ -61,6 +61,83 @@ def load_heightfield_elevation( return data.reshape(header[0], header[1]) +@wp.kernel +def _rasterize_mesh_kernel( + mesh_id: wp.uint64, + x_min: wp.float32, + y_min: wp.float32, + dx: wp.float32, + dy: wp.float32, + z_start: wp.float32, + max_dist: wp.float32, + z_floor: wp.float32, + heights: wp.array2d[wp.float32], +): + row, col = wp.tid() + origin = wp.vec3(x_min + wp.float32(col) * dx, y_min + wp.float32(row) * dy, z_start) + query = wp.mesh_query_ray(mesh_id, origin, wp.vec3(0.0, 0.0, -1.0), max_dist) + heights[row, col] = wp.where(query.result, z_start - query.t, z_floor) + + +def rasterize_mesh_to_heightfield( + mesh: wp.Mesh, + resolution: float, + *, + max_cells_per_axis: int = 4096, +) -> tuple[np.ndarray, tuple[float, float, float, float]]: + """Rasterize a triangle mesh into a heightfield elevation grid. + + Rays are cast straight down onto the mesh on a regular grid spanning the mesh's + XY bounding box. The grid follows Newton's heightfield convention (see + :class:`~newton.Heightfield`): ``heights[row, col]`` samples the surface at + ``x = x_min + col * dx`` (columns map to X) and ``y = y_min + row * dy`` (rows + map to Y). Rays that miss the mesh fall back to the mesh's minimum Z so that + nothing collides in gaps. This is only meaningful for surfaces that are + single-valued in Z (e.g. locomotion terrain). + + Args: + mesh: Triangle mesh to rasterize. Its vertices are taken in their current + frame; transform the mesh into world space beforehand if needed. + resolution: Horizontal grid spacing [m]. Smaller values preserve more + detail at the cost of a larger grid. + max_cells_per_axis: Upper bound on grid rows/columns. If the mesh extent + would exceed this, the effective resolution is coarsened to fit. + + Returns: + A tuple ``(heights, bounds)`` where ``heights`` is a ``(nrow, ncol)`` + float32 array of world-space elevations [m] and ``bounds`` is the mesh's + XY bounding box as ``(x_min, y_min, x_max, y_max)`` [m]. + """ + if resolution <= 0.0: + raise ValueError(f"resolution must be positive, got {resolution}") + if max_cells_per_axis < 2: + raise ValueError(f"max_cells_per_axis must be at least 2, got {max_cells_per_axis}") + + points = mesh.points.numpy() + x_min, y_min, z_min = (float(v) for v in points.min(axis=0)) + x_max, y_max, z_max = (float(v) for v in points.max(axis=0)) + size_x, size_y = x_max - x_min, y_max - y_min + + ncol = max(2, min(int(round(size_x / resolution)) + 1, max_cells_per_axis)) + nrow = max(2, min(int(round(size_y / resolution)) + 1, max_cells_per_axis)) + dx = size_x / (ncol - 1) + dy = size_y / (nrow - 1) + + z_start = z_max + 1.0 + max_dist = (z_start - z_min) + 1.0 + + device = mesh.points.device + heights = wp.empty((nrow, ncol), dtype=wp.float32, device=device) + wp.launch( + _rasterize_mesh_kernel, + dim=(nrow, ncol), + inputs=[mesh.id, x_min, y_min, dx, dy, z_start, max_dist, z_min], + outputs=[heights], + device=device, + ) + return heights.numpy(), (x_min, y_min, x_max, y_max) + + @wp.struct class HeightfieldData: """Per-shape heightfield metadata for collision kernels. @@ -382,192 +459,3 @@ def heightfield_vs_convex_midphase( out_idx = wp.atomic_add(triangle_pairs_count, 0, 1) if out_idx < triangle_pairs.shape[0]: triangle_pairs[out_idx] = wp.vec3i(hfield_shape, other_shape, tri_idx) - - -# Tolerance for rejecting near-parallel rays in the local-space ray intersection. -# Matches raycast.py's PARALLEL_TOL; duplicated here so heightfield queries don't -# depend on raycast.py. -_PARALLEL_TOL = 1e-6 - - -@wp.func -def _ray_intersect_triangle( - ro: wp.vec3, - rd: wp.vec3, - v0: wp.vec3, - v1: wp.vec3, - v2: wp.vec3, -) -> tuple[float, wp.vec3]: - """Moller-Trumbore ray-triangle intersection. - - Returns ``(t, unnormalized_normal)`` on hit, or ``(-1, 0)`` on miss. - Back faces (ray aligned with the face normal) are not culled. - """ - e1 = v1 - v0 - e2 = v2 - v0 - h = wp.cross(rd, e2) - a = wp.dot(e1, h) - if wp.abs(a) < _PARALLEL_TOL: - return -1.0, wp.vec3(0.0) - f = 1.0 / a - s = ro - v0 - u = f * wp.dot(s, h) - if u < 0.0 or u > 1.0: - return -1.0, wp.vec3(0.0) - q = wp.cross(s, e1) - v = f * wp.dot(rd, q) - if v < 0.0 or u + v > 1.0: - return -1.0, wp.vec3(0.0) - t = f * wp.dot(e2, q) - if t < 0.0: - return -1.0, wp.vec3(0.0) - return t, wp.cross(e1, e2) - - -@wp.func -def ray_intersect_heightfield_local( - hfd: HeightfieldData, - elevation_data: wp.array[wp.float32], - ray_origin: wp.vec3, - ray_direction: wp.vec3, -) -> tuple[float, wp.vec3]: - """Ray-heightfield intersection in the heightfield's local frame. - - Slab-clips the ray against the local AABB, then walks the overlapped XY cells - with 2D DDA, testing the two triangles per cell with Moller-Trumbore. Stops - early once the next cell's entry parameter exceeds the best hit found so far. - - Call ``ray_intersect_heightfield`` (in ``geometry.raycast``) from a world-space - kernel -- that thin wrapper does the world-to-local transform and rotates the - returned normal back to world space. - - Args: - hfd: Per-shape heightfield metadata (extents, grid size, z-range, data offset). - elevation_data: Concatenated normalized [0, 1] elevation array. - ray_origin: Ray origin in the heightfield's local frame. - ray_direction: Ray direction in the heightfield's local frame. - - Returns: - The distance along the (local-frame) ray and the unnormalized local-frame - surface normal, or ``-1.0`` and a zero vector on miss. - """ - if hfd.nrow <= 1 or hfd.ncol <= 1: - return -1.0, wp.vec3(0.0) - - ro = ray_origin - rd = ray_direction - - # Slab-clip against the local AABB [-hx, hx] x [-hy, hy] x [min_z, max_z]. - # Explicit float(...) casts make warp treat these as mutable scalars (not constants). - lo = wp.vec3(-hfd.hx, -hfd.hy, hfd.min_z) - hi = wp.vec3(hfd.hx, hfd.hy, hfd.max_z) - t_enter = float(0.0) - t_exit = float(1.0e30) - for i in range(3): - if wp.abs(rd[i]) < _PARALLEL_TOL: - if ro[i] < lo[i] or ro[i] > hi[i]: - return -1.0, wp.vec3(0.0) - else: - inv_d = 1.0 / rd[i] - t1 = (lo[i] - ro[i]) * inv_d - t2 = (hi[i] - ro[i]) * inv_d - t_near = wp.min(t1, t2) - t_far = wp.max(t1, t2) - if t_near > t_enter: - t_enter = t_near - if t_far < t_exit: - t_exit = t_far - if t_enter > t_exit or t_exit < 0.0: - return -1.0, wp.vec3(0.0) - - t_enter = wp.max(t_enter, 0.0) - - dx = 2.0 * hfd.hx / wp.float32(hfd.ncol - 1) - dy = 2.0 * hfd.hy / wp.float32(hfd.nrow - 1) - z_range = hfd.max_z - hfd.min_z - base = hfd.data_offset - - # Starting cell from entry point. - entry = ro + rd * t_enter - col = wp.int32(wp.floor((entry[0] + hfd.hx) / dx)) - row = wp.int32(wp.floor((entry[1] + hfd.hy) / dy)) - col = wp.clamp(col, 0, hfd.ncol - 2) - row = wp.clamp(row, 0, hfd.nrow - 2) - - # DDA deltas and first-boundary parameters per XY axis. - step_col = 0 - t_delta_x = float(1.0e30) - t_next_x = float(1.0e30) - if rd[0] > _PARALLEL_TOL: - step_col = 1 - t_delta_x = dx / rd[0] - t_next_x = (-hfd.hx + wp.float32(col + 1) * dx - ro[0]) / rd[0] - elif rd[0] < -_PARALLEL_TOL: - step_col = -1 - t_delta_x = -dx / rd[0] - t_next_x = (-hfd.hx + wp.float32(col) * dx - ro[0]) / rd[0] - - step_row = 0 - t_delta_y = float(1.0e30) - t_next_y = float(1.0e30) - if rd[1] > _PARALLEL_TOL: - step_row = 1 - t_delta_y = dy / rd[1] - t_next_y = (-hfd.hy + wp.float32(row + 1) * dy - ro[1]) / rd[1] - elif rd[1] < -_PARALLEL_TOL: - step_row = -1 - t_delta_y = -dy / rd[1] - t_next_y = (-hfd.hy + wp.float32(row) * dy - ro[1]) / rd[1] - - best_t = float(1.0e30) - best_normal_local = wp.vec3(0.0, 0.0, 0.0) - - t_cell_enter = t_enter - # A 2D DDA visits at most (nrow + ncol) cells along any straight ray. - max_cells = hfd.nrow + hfd.ncol + 2 - for _ in range(max_cells): - if best_t < t_cell_enter: - break - - x0 = -hfd.hx + wp.float32(col) * dx - y0 = -hfd.hy + wp.float32(row) * dy - x1 = x0 + dx - y1 = y0 + dy - h00 = hfd.min_z + elevation_data[base + row * hfd.ncol + col] * z_range - h10 = hfd.min_z + elevation_data[base + row * hfd.ncol + col + 1] * z_range - h01 = hfd.min_z + elevation_data[base + (row + 1) * hfd.ncol + col] * z_range - h11 = hfd.min_z + elevation_data[base + (row + 1) * hfd.ncol + col + 1] * z_range - - p00 = wp.vec3(x0, y0, h00) - p10 = wp.vec3(x1, y0, h10) - p01 = wp.vec3(x0, y1, h01) - p11 = wp.vec3(x1, y1, h11) - - # Layout: tri 0 = (p00, p10, p11), tri 1 = (p00, p11, p01). - # Matches get_triangle_shape_from_heightfield so collisions and raycasts agree. - t0, n0 = _ray_intersect_triangle(ro, rd, p00, p10, p11) - if t0 >= 0.0 and t0 < best_t: - best_t = t0 - best_normal_local = n0 - t1, n1 = _ray_intersect_triangle(ro, rd, p00, p11, p01) - if t1 >= 0.0 and t1 < best_t: - best_t = t1 - best_normal_local = n1 - - # Step to the next XY cell. - t_cell_enter = wp.min(t_next_x, t_next_y) - if t_cell_enter > t_exit: - break - if t_next_x < t_next_y: - col += step_col - t_next_x += t_delta_x - else: - row += step_row - t_next_y += t_delta_y - if col < 0 or col >= hfd.ncol - 1 or row < 0 or row >= hfd.nrow - 1: - break - - if best_t >= 1.0e30: - return -1.0, wp.vec3(0.0) - - return best_t, best_normal_local diff --git a/newton/_src/utils/import_mjcf.py b/newton/_src/utils/import_mjcf.py index 96f7b91198..d09b9c9e62 100644 --- a/newton/_src/utils/import_mjcf.py +++ b/newton/_src/utils/import_mjcf.py @@ -15,7 +15,7 @@ from ..core import quat_between_axes from ..core.types import Axis, AxisType, Sequence, Transform, vec10 -from ..geometry import Mesh, ShapeFlags +from ..geometry import GeoType, Mesh, ShapeFlags, compute_inertia_shape from ..geometry.types import Heightfield from ..geometry.utils import compute_aabb, compute_inertia_box_mesh from ..sim import JointTargetMode, JointType, ModelBuilder @@ -411,6 +411,15 @@ def parse_mjcf( texture_dir = "." fitaabb = False + inertia_from_geom = compiler_attribs.get("inertiafromgeom", "auto").lower() + if inertia_from_geom not in {"auto", "false", "true"}: + raise ValueError( + f"MJCF compiler inertiafromgeom must be 'auto', 'false', or 'true'; got {inertia_from_geom!r}." + ) + inertia_group_range = tuple(int(value) for value in compiler_attribs.get("inertiagrouprange", "0 5").split()) + if len(inertia_group_range) != 2: + raise ValueError("MJCF compiler inertiagrouprange must contain exactly 2 integers.") + # Parse MJCF compiler and option tags for ONCE and WORLD frequency custom attributes # WORLD frequency attributes use index 0 here; they get remapped during add_world() # Use findall for