diff --git a/source/isaaclab_newton/changelog.d/mzamoramora-collision-decimation.minor.rst b/source/isaaclab_newton/changelog.d/mzamoramora-collision-decimation.minor.rst new file mode 100644 index 000000000000..f2cc59363b0d --- /dev/null +++ b/source/isaaclab_newton/changelog.d/mzamoramora-collision-decimation.minor.rst @@ -0,0 +1,13 @@ +Added +^^^^^ + +* Added :attr:`~isaaclab_newton.physics.NewtonCfg.collision_decimation` to + re-invoke the Newton collision pipeline every ``N`` solver substeps within + a physics tick. Defaults to ``0`` (legacy: one collide per tick). When set + to ``0 < N < num_substeps``, the substep loop in + :meth:`~isaaclab_newton.physics.NewtonManager._run_solver_substeps` calls + the collision pipeline again at the matching substep boundaries so contact + normals reflect the bodies' just-integrated poses. The last substep is + intentionally skipped — its contact set would only feed the next tick. + :meth:`~isaaclab_newton.physics.NewtonCfg.__post_init__` warns when + ``collision_decimation >= num_substeps`` (the gate is silently bypassed). diff --git a/source/isaaclab_newton/changelog.d/mzamoramora-rigid-object-kinematic-fix.rst b/source/isaaclab_newton/changelog.d/mzamoramora-rigid-object-kinematic-fix.rst new file mode 100644 index 000000000000..b858c3c9fafc --- /dev/null +++ b/source/isaaclab_newton/changelog.d/mzamoramora-rigid-object-kinematic-fix.rst @@ -0,0 +1,16 @@ +Fixed +^^^^^ + +* Fixed :class:`~isaaclab_newton.assets.RigidObjectData` crashing with + ``IndexError: tuple index out of range`` for kinematic-enabled + single-body fixed-base rigid objects. The ``is_fixed_base`` branch + in ``_create_simulation_bindings`` indexed ``[:, 0, 0]`` assuming a + 3D ``(count, links, 1)`` layout, but Newton returns a 2D + ``(count, links)`` array when the view contains a single body. + Dispatch on actual ``ndim`` instead so both multi-link fixed-base + articulations and single-body kinematic rigid objects are handled + correctly. Also fixes the matching no-velocity fallback in + ``_create_buffers``: ``_sim_bind_body_com_vel_w`` is now allocated + as ``(num_instances, 1)`` to match the + ``derive_body_acceleration_from_body_com_velocities`` kernel's + 2D signature. diff --git a/source/isaaclab_newton/changelog.d/nut-thread-newton.rst b/source/isaaclab_newton/changelog.d/nut-thread-newton.rst new file mode 100644 index 000000000000..5b5123347f81 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/nut-thread-newton.rst @@ -0,0 +1,6 @@ +Added +^^^^^ + +* Added :attr:`~isaaclab_newton.physics.HydroelasticSDFCfg.moment_matching` + to surface Newton's PhysX-patch-friction analog on the IsaacLab + config side. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py index 43e719d3a580..84d180bb34a8 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py @@ -804,18 +804,22 @@ def _create_simulation_bindings(self) -> None: self._num_bodies = self._root_view.link_count # -- root properties - if self._root_view.is_fixed_base: - self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[ - :, 0, 0 - ] - else: - self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[:, 0] + # Newton's ``get_root_transforms`` / ``get_root_velocities`` return + # either a 2D ``(count, links_per_env)`` array (single-body + # kinematic-enabled rigid objects, e.g. Factory's bolt) or a 3D + # ``(count, links_per_env, 1)`` array (multi-link fixed-base + # articulations). Dispatch on actual ``ndim`` so we don't crash + # with ``IndexError: tuple index out of range`` when a fixed-base + # rigid object yields the 2D layout. + root_xforms = self._root_view.get_root_transforms(SimulationManager.get_state_0()) + self._sim_bind_root_link_pose_w = root_xforms[:, 0, 0] if root_xforms.ndim >= 3 else root_xforms[:, 0] self._sim_bind_root_com_vel_w = self._root_view.get_root_velocities(SimulationManager.get_state_0()) if self._sim_bind_root_com_vel_w is not None: - if self._root_view.is_fixed_base: - self._sim_bind_root_com_vel_w = self._sim_bind_root_com_vel_w[:, 0, 0] - else: - self._sim_bind_root_com_vel_w = self._sim_bind_root_com_vel_w[:, 0] + self._sim_bind_root_com_vel_w = ( + self._sim_bind_root_com_vel_w[:, 0, 0] + if self._sim_bind_root_com_vel_w.ndim >= 3 + else self._sim_bind_root_com_vel_w[:, 0] + ) # -- body properties self._sim_bind_body_com_pos_b = self._root_view.get_attribute("body_com", SimulationManager.get_model())[:, 0] self._sim_bind_body_link_pose_w = self._root_view.get_link_transforms(SimulationManager.get_state_0())[:, 0] @@ -854,11 +858,17 @@ def _create_buffers(self) -> None: "Failed to get root com velocity. If the rigid object is fixed, this is expected. " "Setting root com velocity to zeros." ) + # ``_sim_bind_root_com_vel_w`` is consumed as a 1D array by + # ``get_root_link_vel_from_root_com_vel`` (kernel signature has + # ``com_vel: wp.array(dtype=wp.spatial_vectorf)``), so the fallback + # stays 1D. ``_sim_bind_body_com_vel_w`` feeds the 2D-typed + # ``derive_body_acceleration_from_body_com_velocities`` kernel, so + # the fallback is 2D ``(num_instances, 1)`` to match. self._sim_bind_root_com_vel_w = wp.zeros( (self._num_instances,), dtype=wp.spatial_vectorf, device=self.device ) self._sim_bind_body_com_vel_w = wp.zeros( - (self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + (self._num_instances, 1), dtype=wp.spatial_vectorf, device=self.device ) # -- default root pose and velocity self._default_root_pose = wp.zeros((self._num_instances,), dtype=wp.transformf, device=self.device) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index e99b1eb7abdd..429e59cd390c 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -5,18 +5,32 @@ from __future__ import annotations +import re from collections.abc import Sequence import torch import warp as wp -from newton import ModelBuilder, solvers +from newton import GeoType, ModelBuilder, solvers from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx from pxr import Usd +from isaaclab.physics import PhysicsManager + from isaaclab_newton.physics import NewtonManager +def _compile_sdf_patterns(patterns: list[str]) -> list[re.Pattern]: + """Compile regex patterns with validation, raising on invalid regex.""" + compiled = [] + for i, p in enumerate(patterns): + try: + compiled.append(re.compile(p)) + except re.error as e: + raise ValueError(f"Invalid regex in SDFCfg pattern[{i}]: {p!r} — {e}") from e + return compiled + + def _build_newton_builder_from_mapping( stage: Usd.Stage, sources: Sequence[str], @@ -66,8 +80,6 @@ def _build_newton_builder_from_mapping( # Deformable prim paths are handled by per_world_builder_hooks, not add_usd. # Resolve the regex prim_path patterns to concrete env_0 paths so add_usd # can skip them via ignore_paths. - import re - _deformable_ignore_paths: list[str] = [] if hasattr(NewtonManager, "_deformable_registry"): for entry in NewtonManager._deformable_registry: @@ -81,6 +93,16 @@ def _build_newton_builder_from_mapping( if pat.match(child_path): _deformable_ignore_paths.append(child_path) + # SDF collision requires original triangle meshes for mesh.build_sdf(). + # Convex hull approximation destroys the source geometry, so shapes + # matching SDF patterns must be excluded from approximation here. + # _apply_sdf_config() builds the SDF on each prototype after approximation. + cfg = PhysicsManager._cfg + sdf_cfg = cfg.sdf_cfg if cfg is not None else None # type: ignore[union-attr] + body_pats = _compile_sdf_patterns(sdf_cfg.body_patterns) if sdf_cfg and sdf_cfg.body_patterns else None + shape_pats = _compile_sdf_patterns(sdf_cfg.shape_patterns) if sdf_cfg and sdf_cfg.shape_patterns else None + has_sdf_patterns = body_pats is not None or shape_pats is not None + protos: dict[str, ModelBuilder] = {} for src_path in sources: p = NewtonManager.create_builder(up_axis=up_axis) @@ -94,7 +116,33 @@ def _build_newton_builder_from_mapping( ignore_paths=_deformable_ignore_paths if _deformable_ignore_paths else None, ) if simplify_meshes: - p.approximate_meshes("convex_hull", keep_visual_shapes=True) + if has_sdf_patterns: + sdf_bodies: set[int] = set() + if body_pats is not None: + for bi in range(len(p.body_label)): + if any(pat.search(p.body_label[bi]) for pat in body_pats): + sdf_bodies.add(bi) + + approx_indices = [] + for i in range(len(p.shape_type)): + if p.shape_type[i] != GeoType.MESH: + continue + # Skip shapes that will use SDF (matched by body or shape pattern) + if p.shape_body[i] in sdf_bodies: + continue + if shape_pats is not None: + lbl = p.shape_label[i] if i < len(p.shape_label) else "" + if any(pat.search(lbl) for pat in shape_pats): + continue + approx_indices.append(i) + if approx_indices: + p.approximate_meshes("convex_hull", shape_indices=approx_indices, keep_visual_shapes=True) + else: + p.approximate_meshes("convex_hull", keep_visual_shapes=True) + # Build SDF on prototype before add_builder copies it N times. + # Mesh objects are shared by reference, so SDF is built once and + # all environments inherit it. + NewtonManager._apply_sdf_config(p) protos[src_path] = p # Inject registered sites into prototypes (and global sites into main builder) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi index 176973cd5400..9d832a73d9f1 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi @@ -17,6 +17,7 @@ __all__ = [ "NewtonShapeCfg", "NewtonSolverCfg", "NewtonXPBDManager", + "SDFCfg", "XPBDSolverCfg", ] @@ -26,7 +27,7 @@ from .kamino_manager import NewtonKaminoManager from .kamino_manager_cfg import KaminoSolverCfg from .mjwarp_manager import NewtonMJWarpManager from .mjwarp_manager_cfg import MJWarpSolverCfg -from .newton_collision_cfg import HydroelasticSDFCfg, NewtonCollisionPipelineCfg +from .newton_collision_cfg import HydroelasticSDFCfg, NewtonCollisionPipelineCfg, SDFCfg from .newton_manager import NewtonManager from .newton_manager_cfg import ( NewtonCfg, diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py index 9e75d3153308..72c33403127d 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py @@ -69,6 +69,38 @@ class HydroelasticSDFCfg: Defaults to ``False`` (same as Newton's default). """ + moment_matching: bool = False + """Whether to adjust reduced contact friction so net max moment matches unreduced. + + Only active when ``reduce_contacts`` is True. + + Defaults to ``False`` (same as Newton's default). + """ + + buffer_mult_broad: int = 1 + """Multiplier for preallocated broadphase buffer. + + Defaults to ``1`` (same as Newton's default). + """ + + buffer_mult_iso: int = 1 + """Multiplier for iso-surface extraction buffers. + + Defaults to ``1`` (same as Newton's default). + """ + + buffer_mult_contact: int = 1 + """Multiplier for face contact buffer. + + Defaults to ``1`` (same as Newton's default). + """ + + grid_size: int = 262144 + """Grid size for hydroelastic contact handling (256 * 8 * 128). + + Defaults to ``262144`` (same as Newton's default). + """ + @configclass class NewtonCollisionPipelineCfg: @@ -184,3 +216,119 @@ def to_pipeline_args(self) -> dict[str, Any]: if hydro_cfg is not None: cfg_dict["sdf_hydroelastic_config"] = HydroelasticSDF.Config(**hydro_cfg) return cfg_dict + + +@configclass +class SDFCfg: + """Configuration for SDF mesh collision shapes. + + Specifies how SDF (Signed Distance Field) voxel grids are built and assigned + to bodies or shapes in a Newton model. Bodies and shapes are selected by + regex patterns; the SDF resolution can be set globally or overridden + per-pattern. + + Optional hydroelastic stiffness can be assigned to matched SDF shapes. + Pipeline-level hydroelastic parameters (contact reduction, buffer sizes, + etc.) are configured separately via + :attr:`NewtonCollisionPipelineCfg.sdf_hydroelastic_config`. + + Note: + At least one of :attr:`body_patterns` or :attr:`shape_patterns` must be + set. At least one of :attr:`max_resolution` or + :attr:`target_voxel_size` must be set. + """ + + max_resolution: int | None = None + """Maximum voxel dimension for the SDF grid. + + Must be divisible by 8. Typical values: 128, 256, 512. When both this + and :attr:`target_voxel_size` are set, both are forwarded to Newton's + ``mesh.build_sdf()``; Newton uses :attr:`target_voxel_size` to derive + the actual resolution. + + Defaults to ``None``. + """ + + target_voxel_size: float | None = None + """Target voxel size [m] for the SDF grid. + + When both this and :attr:`max_resolution` are set, both are forwarded to + Newton and this value is used to derive the actual resolution. + + Defaults to ``None``. + """ + + narrow_band_range: tuple[float, float] = (-0.1, 0.1) + """Narrow band distance range (inner, outer) [m]. + + Defines the signed-distance extent stored in the SDF voxel grid. + Negative values are inside the mesh, positive values outside. + + Defaults to ``(-0.1, 0.1)``. + """ + + margin: float | None = None + """Collision margin [m] for SDF shapes. + + When ``None``, the Newton builder default is used. + + Defaults to ``None``. + """ + + body_patterns: list[str] | None = None + """Regex patterns for body labels. + + Matched bodies receive SDF collision shapes on all their mesh geometries. + At least one of :attr:`body_patterns` or :attr:`shape_patterns` must be set. + + Defaults to ``None``. + """ + + shape_patterns: list[str] | None = None + """Regex patterns for shape labels. + + Matched shapes receive SDF collision geometry directly. + At least one of :attr:`body_patterns` or :attr:`shape_patterns` must be set. + + Defaults to ``None``. + """ + + pattern_resolutions: dict[str, int] | None = None + """Per-pattern SDF resolution overrides. + + Maps a regex string to a ``max_resolution`` value. Patterns are evaluated + in insertion order; the first match wins. Unmatched bodies or shapes fall + back to the global :attr:`max_resolution` or :attr:`target_voxel_size`. + + Defaults to ``None``. + """ + + use_visual_meshes: bool = False + """Whether to create collision shapes from visual meshes. + + When ``True``, matched bodies that lack explicit collision geometry have SDF + collision shapes built from their visual meshes instead. + + Defaults to ``False``. + """ + + k_hydro: float | None = None + """Hydroelastic stiffness [Pa] assigned to matched SDF shapes. + + When ``None``, no ``HYDROELASTIC`` flag is set and hydroelastic contacts are + disabled for these shapes. When set, matched shapes receive the flag with + this stiffness value. + + Defaults to ``None``. + """ + + hydroelastic_shape_patterns: list[str] | None = None + """Regex patterns restricting which SDF shapes receive hydroelastic stiffness. + + Only relevant when :attr:`k_hydro` is set. When ``None``, all SDF shapes + matched by :attr:`body_patterns` or :attr:`shape_patterns` get the + hydroelastic flag. When set, only shapes whose labels match at least one + pattern here receive it. + + Defaults to ``None``. + """ diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 500d07b726e7..77d3fc6c7835 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -10,10 +10,12 @@ import contextlib import ctypes import logging +import re from abc import abstractmethod from collections.abc import Callable from typing import TYPE_CHECKING +import torch import warp as wp # Load CUDA runtime for relaxed-mode graph capture (RTX-compatible). @@ -199,6 +201,7 @@ class NewtonManager(PhysicsManager): _solver_dt: float = 1.0 / 200.0 _num_substeps: int = 1 _decimation: int = 1 + _collision_decimation: int = 0 _num_envs: int | None = None # Newton model and state @@ -230,6 +233,10 @@ class NewtonManager(PhysicsManager): _world_reset_mask: wp.array | None = None # (num_envs,) wp.int32 — for SolverKamino.reset(world_mask=...) _fk_reset_mask: wp.array | None = None # (articulation_count,) wp.bool — for eval_fk(mask=...) + # Per-world mask of NaN-divergent envs queued for sanitize at next reset. + # Populated by :meth:`sanitize_nan_envs`, drained by :meth:`sanitize_pending_nan_envs`. + _nan_env_mask_pending_reset: torch.Tensor | None = None + # Newton actuator adapter (owns actuators and double-buffered states) _adapter: NewtonActuatorAdapter | None = None # In-graph hooks invoked after the actuator step and before the solver @@ -677,6 +684,7 @@ def clear(cls): # Per-world reset masks NewtonManager._world_reset_mask = None NewtonManager._fk_reset_mask = None + NewtonManager._nan_env_mask_pending_reset = None NewtonManager._graph = None NewtonManager._graph_capture_pending = False NewtonManager._newton_stage_path = None @@ -891,6 +899,212 @@ def add_model_change(cls, change: SolverNotifyFlags) -> None: """Register a model change to notify the solver.""" cls._model_changes.add(change) + @staticmethod + def _build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, label: str) -> bool: + """Build SDF on a mesh, resolving per-pattern resolution overrides. + + Args: + mesh: Newton mesh object to build SDF on. + sdf_cfg: The active :class:`SDFCfg` instance. + res_overrides: Compiled ``(pattern, resolution)`` pairs, or ``None``. + label: Shape label used for pattern resolution matching. + + Returns: + ``True`` if SDF was built, ``False`` if skipped (no mesh source). + """ + if mesh is None: + logger.warning(f"SDF: shape '{label}' matched but has no mesh source. Skipping SDF build.") + return False + if mesh.sdf is not None: + mesh.clear_sdf() + resolution = sdf_cfg.max_resolution + if res_overrides is not None: + for pat, res in res_overrides: + if pat.search(label): + resolution = res + break + sdf_kwargs: dict = dict(narrow_band_range=sdf_cfg.narrow_band_range) + if resolution is not None: + sdf_kwargs["max_resolution"] = resolution + if sdf_cfg.target_voxel_size is not None: + sdf_kwargs["target_voxel_size"] = sdf_cfg.target_voxel_size + mesh.build_sdf(**sdf_kwargs) + return True + + @classmethod + def _create_sdf_collision_from_visual( + cls, builder: ModelBuilder, sdf_shape_indices: set[int], sdf_cfg, res_overrides, hydro_patterns=None + ): + """Create collision shapes from visual meshes for matched bodies lacking collision geometry. + + Args: + builder: Newton model builder to modify. + sdf_shape_indices: Shape indices that matched SDF patterns. + sdf_cfg: The active :class:`SDFCfg` instance. + res_overrides: Compiled ``(pattern, resolution)`` pairs, or ``None``. + hydro_patterns: Compiled hydroelastic shape patterns, or ``None`` + (meaning all shapes get hydroelastic if ``k_hydro`` is set). + + Returns: + Tuple of ``(num_added, num_hydro)`` counts. + """ + from newton import ShapeFlags + + matched_bodies: set[int] = {builder.shape_body[si] for si in sdf_shape_indices} + bodies_with_collision: set[int] = set() + for si in range(builder.shape_count): + if builder.shape_flags[si] & ShapeFlags.COLLIDE_SHAPES and builder.shape_body[si] in matched_bodies: + bodies_with_collision.add(builder.shape_body[si]) + + num_added = 0 + num_hydro = 0 + for body_idx in matched_bodies - bodies_with_collision: + visual_si = None + for si in sdf_shape_indices: + if builder.shape_body[si] == body_idx and builder.shape_source[si] is not None: + visual_si = si + break + if visual_si is None: + body_lbl = builder.body_label[body_idx] + logger.warning(f"SDF: body '{body_lbl}' matched but has no visual mesh to create collision from.") + continue + + mesh = builder.shape_source[visual_si] + cls._build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, builder.shape_label[visual_si]) + + shape_lbl = builder.shape_label[visual_si] + enable_hydro = False + if sdf_cfg.k_hydro is not None: + enable_hydro = hydro_patterns is None or any(p.search(shape_lbl) for p in hydro_patterns) + + shape_cfg_kwargs: dict = dict( + density=0.0, + has_shape_collision=True, + has_particle_collision=True, + is_visible=False, + ) + if sdf_cfg.margin is not None: + shape_cfg_kwargs["margin"] = sdf_cfg.margin + if enable_hydro: + shape_cfg_kwargs["is_hydroelastic"] = True + shape_cfg_kwargs["kh"] = sdf_cfg.k_hydro + + body_lbl = builder.body_label[body_idx] + builder.add_shape_mesh( + body=body_idx, + xform=builder.shape_transform[visual_si], + mesh=mesh, + scale=builder.shape_scale[visual_si], + cfg=ModelBuilder.ShapeConfig(**shape_cfg_kwargs), + label=f"{body_lbl}/sdf_collision", + ) + num_added += 1 + if enable_hydro: + num_hydro += 1 + + return num_added, num_hydro + + @classmethod + def _apply_sdf_config(cls, builder: ModelBuilder): + """Apply SDF collision and optional hydroelastic flags to matching mesh shapes. + + Reads :class:`SDFCfg` from the active physics config. Collects shapes + matching body/shape regex patterns, builds SDF on their meshes, and + optionally sets the ``HYDROELASTIC`` flag with :attr:`SDFCfg.k_hydro`. + + Args: + builder: Newton model builder to modify (before finalization). + """ + from newton import GeoType, ShapeFlags + + cfg = PhysicsManager._cfg + if cfg is None: + return + sdf_cfg = getattr(cfg, "sdf_cfg", None) + if sdf_cfg is None: + return + + if sdf_cfg.max_resolution is None and sdf_cfg.target_voxel_size is None: + logger.warning("SDFCfg provided but neither max_resolution nor target_voxel_size is set. SDF disabled.") + return + + def _compile(patterns: list[str] | None, field: str) -> list[re.Pattern] | None: + if not patterns: + return None + compiled = [] + for i, p in enumerate(patterns): + try: + compiled.append(re.compile(p)) + except re.error as e: + raise ValueError(f"Invalid regex in SDFCfg.{field}[{i}]: {p!r} — {e}") from e + return compiled + + body_patterns = _compile(sdf_cfg.body_patterns, "body_patterns") + shape_patterns = _compile(sdf_cfg.shape_patterns, "shape_patterns") + res_overrides = None + if sdf_cfg.pattern_resolutions: + res_overrides = [] + for p, r in sdf_cfg.pattern_resolutions.items(): + try: + res_overrides.append((re.compile(p), r)) + except re.error as e: + raise ValueError(f"Invalid regex in SDFCfg.pattern_resolutions key {p!r} — {e}") from e + hydro_patterns = None + if sdf_cfg.k_hydro is not None: + hydro_patterns = _compile(sdf_cfg.hydroelastic_shape_patterns, "hydroelastic_shape_patterns") + + if body_patterns is None and shape_patterns is None: + logger.warning("SDFCfg has no body_patterns or shape_patterns set. No shapes will receive SDF.") + return + + body_to_shapes: dict[int, list[int]] = {} + for si in range(builder.shape_count): + if builder.shape_type[si] == GeoType.MESH: + body_to_shapes.setdefault(builder.shape_body[si], []).append(si) + + sdf_shape_indices: set[int] = set() + + if body_patterns is not None: + for body_idx in range(len(builder.body_label)): + if any(p.search(builder.body_label[body_idx]) for p in body_patterns): + sdf_shape_indices.update(body_to_shapes.get(body_idx, [])) + + if shape_patterns is not None: + for shape_indices in body_to_shapes.values(): + for si in shape_indices: + if any(p.search(builder.shape_label[si]) for p in shape_patterns): + sdf_shape_indices.add(si) + + num_patched = 0 + num_hydro = 0 + for si in sdf_shape_indices: + if not (builder.shape_flags[si] & ShapeFlags.COLLIDE_SHAPES): + continue + if not cls._build_sdf_on_mesh(builder.shape_source[si], sdf_cfg, res_overrides, builder.shape_label[si]): + continue + if sdf_cfg.margin is not None: + builder.shape_margin[si] = sdf_cfg.margin + if sdf_cfg.k_hydro is not None: + apply_hydro = hydro_patterns is None or any(p.search(builder.shape_label[si]) for p in hydro_patterns) + if apply_hydro: + builder.shape_flags[si] |= ShapeFlags.HYDROELASTIC + builder.shape_material_kh[si] = sdf_cfg.k_hydro + num_hydro += 1 + num_patched += 1 + + num_added = 0 + if sdf_cfg.use_visual_meshes: + num_added, hydro_from_visual = cls._create_sdf_collision_from_visual( + builder, sdf_shape_indices, sdf_cfg, res_overrides, hydro_patterns + ) + num_hydro += hydro_from_visual + + hydro_msg = f", {num_hydro} hydroelastic shape(s)" if sdf_cfg.k_hydro is not None else "" + logger.info( + f"SDF config: {num_added} collision shape(s) added, {num_patched} existing shape(s) patched{hydro_msg}. " + f"(max_resolution={sdf_cfg.max_resolution}, narrow_band={sdf_cfg.narrow_band_range})" + ) + @classmethod def invalidate_fk( cls, @@ -1218,6 +1432,7 @@ def initialize_solver(cls) -> None: with Timer(name="newton_initialize_solver", msg="Initialize solver took:"): NewtonManager._num_substeps = cfg.num_substeps # type: ignore[union-attr] + NewtonManager._collision_decimation = cfg.collision_decimation # type: ignore[union-attr] NewtonManager._solver_dt = cls.get_physics_dt() / cls._num_substeps NewtonManager._collision_cfg = cfg.collision_cfg # type: ignore[union-attr] @@ -1439,10 +1654,17 @@ def _capture_relaxed_graph(cls, device: str): @classmethod def _run_solver_substeps(cls, contacts) -> None: """Run ``num_substeps`` solver iterations, handling double-buffered state swap.""" + collide_every = cls._collision_decimation + # Last substep is skipped: its contact set would only feed the next tick's + # top-of-loop collide(), not this one. + collide_mid_loop = collide_every > 0 and cls._needs_collision_pipeline and contacts is not None + if cls._use_single_state: - for _ in range(cls._num_substeps): + for i in range(cls._num_substeps): cls._step_solver(cls._state_0, cls._state_0, cls._control, contacts, cls._solver_dt) cls._state_0.clear_forces() + if collide_mid_loop and (i + 1) % collide_every == 0 and i + 1 < cls._num_substeps: + cls._collision_pipeline.collide(cls._state_0, contacts) else: cfg = PhysicsManager._cfg need_copy_on_last = (cfg is not None and cfg.use_cuda_graph) and cls._num_substeps % 2 == 1 # type: ignore[union-attr] @@ -1453,6 +1675,8 @@ def _run_solver_substeps(cls, contacts) -> None: else: NewtonManager._state_0, NewtonManager._state_1 = cls._state_1, cls._state_0 cls._state_0.clear_forces() + if collide_mid_loop and (i + 1) % collide_every == 0 and i + 1 < cls._num_substeps: + cls._collision_pipeline.collide(cls._state_0, contacts) @classmethod def _update_sensors(cls, contacts) -> None: @@ -1469,6 +1693,141 @@ def _update_sensors(cls, contacts) -> None: for sensor in cls._newton_contact_sensors.values(): sensor.update(cls._state_0, eval_contacts) + # ------------------------------------------------------------------ + # NaN recovery + # ------------------------------------------------------------------ + + @classmethod + def sanitize_world_state(cls, env_ids) -> None: + """Reset solver internals + Newton State buffers for the given worlds. + + Used by tasks that detect NaN-divergent worlds and want to recover + them in place (vs throwing the episode away). Zeroes per-world + MJWarp solver scratch buffers (``qacc_warmstart``, ``qfrc_*``, + ``cacc``, ``cfrc_*``) and Newton State velocity/force buffers + (``joint_qd``, ``body_qd``, ``body_f``, ``body_qdd``, + ``body_parent_f``) at the indexed worlds, then runs + :func:`newton.eval_fk` to re-derive ``state.body_q`` from + ``joint_q``. + + ``mjwarp.reset_data`` covers ``qacc_warmstart`` and the contact + arrays but leaves the derived ``qfrc_*`` family and the COM-frame + ``cfrc_int``/``cacc``/``cfrc_ext`` untouched. Empirically those + latter buffers re-divergence the same world on the next solver + step (``cfrc_int`` is read-modify-written via ``wp.atomic_add`` in + ``mujoco_warp/_src/smooth.py:_cfrc_backward``, so a stale NaN + survives the next ``rne()`` call). This implementation walks all + 17 named MJWarp fields explicitly to close that gap. See + ``newton-physics/newton#1266`` for the upstream discussion. + + No-op for solvers without a ``mjw_data`` attribute (XPBD, + Featherstone, Kamino) — those solvers do not exhibit the MuJoCo + warm-start contamination pattern this addresses. + + Args: + env_ids: 1-D ``int`` torch tensor of world IDs to sanitize, on + the simulation device. + """ + import newton as _newton # noqa: PLC0415 + + if cls._model is None or cls._solver is None or cls._num_envs is None: + return + if env_ids.numel() == 0: + return + + mjw_data = getattr(cls._solver, "mjw_data", None) + if mjw_data is None: + return # not the MJWarp backend; nothing to scrub + + env_ids_list = env_ids.tolist() + for field_name in ( + "qacc_warmstart", + "qacc", + "qacc_smooth", + "qfrc_applied", + "qfrc_bias", + "qfrc_spring", + "qfrc_damper", + "qfrc_gravcomp", + "qfrc_fluid", + "qfrc_passive", + "qfrc_actuator", + "qfrc_smooth", + "qfrc_constraint", + "qfrc_inverse", + "cacc", + "cfrc_int", + "cfrc_ext", + ): + arr = getattr(mjw_data, field_name, None) + if arr is None: + continue + t = wp.to_torch(arr) + for env_id in env_ids_list: + t[env_id] = 0.0 + + state = cls._state_0 + model = cls._model + num_envs = cls._num_envs + nd_per_env = model.joint_dof_count // num_envs + nb_per_env = model.body_count // num_envs + for buf_name, per_env in ( + ("joint_qd", nd_per_env), + ("body_qd", nb_per_env), + ("body_f", nb_per_env), + ("body_qdd", nb_per_env), + ("body_parent_f", nb_per_env), + ): + buf = getattr(state, buf_name, None) + if buf is None: + continue + wp.to_torch(buf).view(num_envs, per_env, -1)[env_ids] = 0.0 + + _newton.eval_fk(model, state.joint_q, state.joint_qd, state) + body_q_prev = getattr(state, "body_q_prev", None) + if body_q_prev is not None: + bq = wp.to_torch(state.body_q).view(num_envs, nb_per_env, -1) + wp.to_torch(body_q_prev).view(num_envs, nb_per_env, -1)[env_ids] = bq[env_ids] + + @classmethod + def flag_nan_envs(cls, mask: torch.Tensor) -> None: + """Flag NaN-divergent envs for sanitization at the next episode reset. + + Call from a task's per-step NaN detection (e.g. ``torch.isnan`` over + obs/state tensors). Pure bookkeeping: the mask is ORed into + :attr:`_nan_env_mask_pending_reset`; no solver state is touched yet. + :meth:`sanitize_pending_nan_envs` drains the queue at the next reset. + + Tasks that need to keep training between detection and reset should + ``torch.nan_to_num`` their obs/state/reward themselves — flagging is + the only side effect here. + + No-op when ``mask.any()`` is False. + + Args: + mask: 1-D ``bool`` torch tensor of shape ``(num_worlds,)``; True for + envs just detected as NaN-divergent. + """ + if not mask.any(): + return + if cls._nan_env_mask_pending_reset is None: + cls._nan_env_mask_pending_reset = torch.zeros_like(mask) + cls._nan_env_mask_pending_reset |= mask + + @classmethod + def sanitize_pending_nan_envs(cls) -> None: + """Drain the pending-reset NaN queue and sanitize those envs. + + Call from a task's reset hook (e.g. ``_reset_idx``) before re-init. + No-op when no envs were queued via :meth:`flag_nan_envs` since the + last call. + """ + if cls._nan_env_mask_pending_reset is None or not cls._nan_env_mask_pending_reset.any(): + return + nan_ids = cls._nan_env_mask_pending_reset.nonzero(as_tuple=False).squeeze(-1) + cls.sanitize_world_state(nan_ids) + cls._nan_env_mask_pending_reset.zero_() + # ------------------------------------------------------------------ # Composite stepping routines # ------------------------------------------------------------------ diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py index 6ff646aff57b..c2d6aabda565 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py @@ -7,16 +7,19 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING from isaaclab.physics import PhysicsCfg from isaaclab.utils.configclass import configclass -from .newton_collision_cfg import NewtonCollisionPipelineCfg +from .newton_collision_cfg import NewtonCollisionPipelineCfg, SDFCfg if TYPE_CHECKING: from isaaclab_newton.physics import NewtonManager +logger = logging.getLogger(__name__) + @configclass class NewtonSolverCfg: @@ -97,6 +100,9 @@ class NewtonCfg(PhysicsCfg): num_substeps: int = 1 """Number of substeps to use for the solver.""" + collision_decimation: int = 0 + """Re-collide every N solver substeps within a physics tick (``0`` = once per tick).""" + debug_mode: bool = False """Whether to enable debug mode for the solver.""" @@ -138,6 +144,15 @@ class NewtonCfg(PhysicsCfg): :class:`NewtonShapeCfg` for the declared fields. """ + sdf_cfg: SDFCfg | None = None + """SDF mesh collision configuration. + + When set, matched bodies/shapes receive SDF collision and optional + hydroelastic stiffness. See :class:`~isaaclab_newton.physics.SDFCfg` + for the declared fields. When ``None`` (default), no SDF collision is + configured. + """ + def __post_init__(self): # NewtonCfg.class_type is auto-derived from solver_cfg.class_type. # Refuse a user-set value: setting both is ambiguous and was @@ -149,3 +164,12 @@ def __post_init__(self): self.solver_cfg = MJWarpSolverCfg() self.class_type = self.solver_cfg.class_type + + # Mid-tick re-collide is silently disabled when collision_decimation >= num_substeps. + if self.collision_decimation > 0 and self.collision_decimation >= self.num_substeps: + logger.warning( + "NewtonCfg.collision_decimation=%d is >= num_substeps=%d; mid-tick re-collide is disabled. " + "Set 0 < collision_decimation < num_substeps to enable.", + self.collision_decimation, + self.num_substeps, + ) diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py index f8ac77586cc0..8458f466632c 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -134,6 +134,29 @@ def test_newton_cfg_post_init_propagates_class_type( assert cfg.class_type.__name__ == expected_manager.__name__ +@pytest.mark.parametrize( + "num_substeps, collision_decimation, should_warn", + [ + (8, 0, False), # Default: feature disabled, no warning. + (8, 1, False), # Valid: re-collide every substep. + (8, 2, False), # Valid: re-collide every 2 substeps. + (8, 7, False), # Valid edge: one mid-loop re-collide at i=6. + (8, 8, True), # Equal to num_substeps: gate never fires. + (8, 16, True), # Larger than num_substeps: gate never fires. + ], +) +def test_newton_cfg_collision_decimation_warning(num_substeps, collision_decimation, should_warn, caplog): + """``NewtonCfg.__post_init__`` warns when ``collision_decimation >= num_substeps``.""" + import logging + + with caplog.at_level(logging.WARNING, logger="isaaclab_newton.physics.newton_manager_cfg"): + cfg = NewtonCfg(num_substeps=num_substeps, collision_decimation=collision_decimation) + warned = any("collision_decimation" in rec.getMessage() for rec in caplog.records) + assert warned is should_warn + # Cfg field round-trips regardless of warning. + assert cfg.collision_decimation == collision_decimation + + # --------------------------------------------------------------------------- # Manager class hierarchy and factory contracts # --------------------------------------------------------------------------- @@ -272,3 +295,68 @@ def test_mjwarp_internal_contacts_with_collision_cfg_raises(): with pytest.raises(ValueError, match="collision_cfg cannot be set"): sim.reset() + + +@pytest.mark.parametrize( + "num_substeps, collision_decimation, expected_mid_loop_collides", + [ + (8, 0, 0), # Feature disabled. + (8, 2, 3), # Re-collide after substeps 2, 4, 6 (skip last). + (8, 4, 1), # Re-collide after substep 4 only. + (8, 7, 1), # Re-collide after substep 7 only. + (8, 8, 0), # Gated off (>= num_substeps). + ], +) +def test_collision_decimation_invokes_mid_loop_collide(num_substeps, collision_decimation, expected_mid_loop_collides): + """``_run_solver_substeps`` re-invokes ``collide`` at the expected substeps. + + Wraps :attr:`NewtonManager._collision_pipeline.collide` with a counter and + runs one physics tick. The collide-call count is ``1`` (top-of-tick) plus + one per matching mid-loop substep, excluding the last substep. + + The scene has a free-joint sphere falling onto a ground plane so the + broadphase actually generates pairs — guards against a future change + that skips ``collide()`` when there are no collidable shapes. + """ + sim_cfg = SimulationCfg( + dt=1.0 / 120.0, + device="cuda:0", + gravity=(0.0, 0.0, -9.81), + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg(use_mujoco_contacts=False), + num_substeps=num_substeps, + collision_decimation=collision_decimation, + use_cuda_graph=False, + ), + ) + + with build_simulation_context(sim_cfg=sim_cfg) as sim: + builder = NewtonManager.create_builder() + body = builder.add_body(mass=1.0) + builder.add_joint_free(child=body) + builder.add_shape_sphere(body=body, radius=0.05) + builder.add_ground_plane() + # Lift the sphere to 0.5 m above the plane so the scene is non-degenerate. + # joint_q for a free joint is [tx, ty, tz, qx, qy, qz, qw]. + builder.joint_q[-7:] = [0.0, 0.0, 0.5, 0.0, 0.0, 0.0, 1.0] + NewtonManager.set_builder(builder) + sim.reset() + + # Wrap collide() with a counter — must run after sim.reset() so the + # pipeline is allocated, and use_cuda_graph=False so the wrapped + # Python callable isn't bypassed by a captured graph. + calls = {"n": 0} + original_collide = NewtonManager._collision_pipeline.collide + + def counting_collide(state, contacts): + calls["n"] += 1 + return original_collide(state, contacts) + + NewtonManager._collision_pipeline.collide = counting_collide + try: + sim.step(render=False) + finally: + NewtonManager._collision_pipeline.collide = original_collide + + # Expect: 1 (top-of-tick) + expected_mid_loop_collides. + assert calls["n"] == 1 + expected_mid_loop_collides diff --git a/source/isaaclab_newton/test/physics/test_sdf_config.py b/source/isaaclab_newton/test/physics/test_sdf_config.py new file mode 100644 index 000000000000..b2ed547753e8 --- /dev/null +++ b/source/isaaclab_newton/test/physics/test_sdf_config.py @@ -0,0 +1,487 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for SDF collision configuration and application logic.""" + +import re +from unittest.mock import MagicMock, patch + +from newton import GeoType, ModelBuilder, ShapeFlags + + +class TestBuildSdfOnMesh: + """Tests for NewtonManager._build_sdf_on_mesh.""" + + @staticmethod + def _make_sdf_cfg(max_resolution=256, narrow_band_range=(-0.1, 0.1), target_voxel_size=None): + cfg = MagicMock() + cfg.max_resolution = max_resolution + cfg.narrow_band_range = narrow_band_range + cfg.target_voxel_size = target_voxel_size + return cfg + + def test_none_mesh_is_noop(self): + """Passing None as mesh should not raise.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + sdf_cfg = self._make_sdf_cfg() + NewtonManager._build_sdf_on_mesh(None, sdf_cfg, None, "test_label") + + def test_builds_sdf_with_max_resolution(self): + """SDF is built on mesh with max_resolution and narrow_band_range.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=128) + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, None, "test_label") + + mesh.build_sdf.assert_called_once_with(narrow_band_range=(-0.1, 0.1), max_resolution=128) + + def test_clears_existing_sdf_before_rebuild(self): + """Existing SDF on mesh is cleared before building a new one.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = "existing_sdf" + sdf_cfg = self._make_sdf_cfg() + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, None, "test_label") + + mesh.clear_sdf.assert_called_once() + mesh.build_sdf.assert_called_once() + + def test_target_voxel_size_passed_alongside_resolution(self): + """When target_voxel_size is set, it is passed alongside max_resolution.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=256, target_voxel_size=0.005) + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, None, "test_label") + + call_kwargs = mesh.build_sdf.call_args[1] + assert call_kwargs["target_voxel_size"] == 0.005 + assert call_kwargs["max_resolution"] == 256 + + def test_resolution_override_by_pattern(self): + """Per-pattern resolution override is applied when label matches.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=256) + res_overrides = [(re.compile(".*elbow.*"), 128)] + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, "/World/Robot/elbow_link/collision") + + call_kwargs = mesh.build_sdf.call_args[1] + assert call_kwargs["max_resolution"] == 128 + + def test_resolution_override_no_match_uses_global(self): + """When label doesn't match any override, global max_resolution is used.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=256) + res_overrides = [(re.compile(".*elbow.*"), 128)] + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, "/World/Robot/wrist_link/collision") + + call_kwargs = mesh.build_sdf.call_args[1] + assert call_kwargs["max_resolution"] == 256 + + def test_resolution_override_first_match_wins(self): + """First matching pattern in res_overrides determines resolution.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=256) + res_overrides = [ + (re.compile(".*link.*"), 64), + (re.compile(".*elbow.*"), 128), + ] + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, "/World/Robot/elbow_link/collision") + + call_kwargs = mesh.build_sdf.call_args[1] + assert call_kwargs["max_resolution"] == 64 # ".*link.*" matches first + + +class TestApplySdfConfig: + """Tests for NewtonManager._apply_sdf_config shape index collection and patching.""" + + @staticmethod + def _make_builder(bodies, shapes): + """Create a minimal ModelBuilder-like mock. + + Args: + bodies: List of body label strings. + shapes: List of dicts with keys: body_idx, label, geo_type, flags, source. + """ + builder = MagicMock(spec=ModelBuilder) + builder.body_label = bodies + builder.shape_count = len(shapes) + builder.shape_type = [s["geo_type"] for s in shapes] + builder.shape_body = [s["body_idx"] for s in shapes] + builder.shape_label = [s["label"] for s in shapes] + builder.shape_flags = [s["flags"] for s in shapes] + builder.shape_source = [s.get("source") for s in shapes] + builder.shape_margin = [0.0] * len(shapes) + builder.shape_material_kh = [0.0] * len(shapes) + return builder + + @staticmethod + def _make_cfg( + body_patterns=None, + shape_patterns=None, + max_resolution=256, + k_hydro=None, + hydroelastic_shape_patterns=None, + ): + cfg = MagicMock() + cfg.sdf_cfg = MagicMock() + cfg.sdf_cfg.max_resolution = max_resolution + cfg.sdf_cfg.target_voxel_size = None + cfg.sdf_cfg.narrow_band_range = (-0.1, 0.1) + cfg.sdf_cfg.margin = None + cfg.sdf_cfg.body_patterns = body_patterns + cfg.sdf_cfg.shape_patterns = shape_patterns + cfg.sdf_cfg.pattern_resolutions = None + cfg.sdf_cfg.use_visual_meshes = False + cfg.sdf_cfg.k_hydro = k_hydro + cfg.sdf_cfg.hydroelastic_shape_patterns = hydroelastic_shape_patterns + return cfg + + def test_no_sdf_cfg_is_noop(self): + """_apply_sdf_config returns early when sdf_cfg is None.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = MagicMock(spec=ModelBuilder) + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as pm: + pm._cfg = MagicMock() + pm._cfg.sdf_cfg = None + NewtonManager._apply_sdf_config(builder) + assert not builder.method_calls + + def test_no_patterns_warns(self): + """_apply_sdf_config warns when no patterns are set.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = MagicMock(spec=ModelBuilder) + cfg = self._make_cfg(body_patterns=None, shape_patterns=None) + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as pm: + pm._cfg = cfg + with patch("isaaclab_newton.physics.newton_manager.logger") as mock_logger: + NewtonManager._apply_sdf_config(builder) + mock_logger.warning.assert_called() + + def test_body_pattern_collects_shapes(self): + """Shapes under matching bodies are collected for SDF.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/elbow", "/World/Robot/wrist"] + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/elbow/col", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": MagicMock(sdf=None), + }, + { + "body_idx": 1, + "label": "/World/Robot/wrist/col", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": MagicMock(sdf=None), + }, + ] + builder = self._make_builder(bodies, shapes) + cfg = self._make_cfg(body_patterns=[".*elbow.*"]) + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as pm: + pm._cfg = cfg + NewtonManager._apply_sdf_config(builder) + + shapes[0]["source"].build_sdf.assert_called_once() + shapes[1]["source"].build_sdf.assert_not_called() + + def test_hydroelastic_flag_set_when_k_hydro(self): + """HYDROELASTIC flag is set on matched shapes when k_hydro is provided.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/elbow"] + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/elbow/col", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": MagicMock(sdf=None), + }, + ] + builder = self._make_builder(bodies, shapes) + cfg = self._make_cfg(body_patterns=[".*elbow.*"], k_hydro=1e10) + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as pm: + pm._cfg = cfg + NewtonManager._apply_sdf_config(builder) + + assert builder.shape_flags[0] & ShapeFlags.HYDROELASTIC + assert builder.shape_material_kh[0] == 1e10 + + def test_hydroelastic_shape_patterns_filter(self): + """hydroelastic_shape_patterns limits which shapes get HYDROELASTIC flag.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/elbow", "/World/Robot/wrist"] + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/elbow/col", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": MagicMock(sdf=None), + }, + { + "body_idx": 1, + "label": "/World/Robot/wrist/col", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": MagicMock(sdf=None), + }, + ] + builder = self._make_builder(bodies, shapes) + cfg = self._make_cfg( + body_patterns=[".*"], + k_hydro=1e10, + hydroelastic_shape_patterns=[".*elbow.*"], + ) + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as pm: + pm._cfg = cfg + NewtonManager._apply_sdf_config(builder) + + shapes[0]["source"].build_sdf.assert_called_once() + shapes[1]["source"].build_sdf.assert_called_once() + assert builder.shape_flags[0] & ShapeFlags.HYDROELASTIC + assert not (builder.shape_flags[1] & ShapeFlags.HYDROELASTIC) + + def test_shape_pattern_matching(self): + """shape_patterns directly matches shape labels.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/body"] + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/body/Gear_col", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": MagicMock(sdf=None), + }, + { + "body_idx": 0, + "label": "/World/Robot/body/frame_col", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": MagicMock(sdf=None), + }, + ] + builder = self._make_builder(bodies, shapes) + cfg = self._make_cfg(shape_patterns=[".*Gear.*"]) + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as pm: + pm._cfg = cfg + NewtonManager._apply_sdf_config(builder) + + shapes[0]["source"].build_sdf.assert_called_once() + shapes[1]["source"].build_sdf.assert_not_called() + + def test_non_mesh_shapes_skipped(self): + """Non-mesh shapes are never collected for SDF.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/elbow"] + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/elbow/box", + "geo_type": GeoType.BOX, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": None, + }, + ] + builder = self._make_builder(bodies, shapes) + cfg = self._make_cfg(body_patterns=[".*elbow.*"]) + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as pm: + pm._cfg = cfg + NewtonManager._apply_sdf_config(builder) + + +class TestCreateSdfCollisionFromVisual: + """Tests for NewtonManager._create_sdf_collision_from_visual.""" + + @staticmethod + def _make_builder(bodies, shapes): + """Create a minimal ModelBuilder-like mock.""" + builder = MagicMock(spec=ModelBuilder) + builder.body_label = bodies + builder.shape_count = len(shapes) + builder.shape_type = [s["geo_type"] for s in shapes] + builder.shape_body = [s["body_idx"] for s in shapes] + builder.shape_label = [s["label"] for s in shapes] + builder.shape_flags = [s["flags"] for s in shapes] + builder.shape_source = [s.get("source") for s in shapes] + builder.shape_transform = [s.get("xform", (0, 0, 0, 0, 0, 0, 1)) for s in shapes] + builder.shape_scale = [s.get("scale", (1, 1, 1)) for s in shapes] + builder.shape_margin = [0.0] * len(shapes) + builder.shape_material_kh = [0.0] * len(shapes) + return builder + + @staticmethod + def _make_sdf_cfg(k_hydro=None, margin=None): + cfg = MagicMock() + cfg.max_resolution = 256 + cfg.target_voxel_size = None + cfg.narrow_band_range = (-0.1, 0.1) + cfg.margin = margin + cfg.k_hydro = k_hydro + return cfg + + def test_creates_collision_from_visual_mesh(self): + """A body with only a visual mesh (no COLLIDE_SHAPES) gets a new collision shape.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/arm"] + mesh = MagicMock(sdf=None) + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/arm/visual", + "geo_type": GeoType.MESH, + "flags": 0, + "source": mesh, + }, + ] + builder = self._make_builder(bodies, shapes) + sdf_cfg = self._make_sdf_cfg() + + num_added, num_hydro = NewtonManager._create_sdf_collision_from_visual(builder, {0}, sdf_cfg, None) + + assert num_added == 1 + assert num_hydro == 0 + builder.add_shape_mesh.assert_called_once() + mesh.build_sdf.assert_called_once() + + def test_skips_body_with_existing_collision(self): + """A body that already has a COLLIDE_SHAPES shape is not given a visual-mesh collision.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/arm"] + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/arm/collision", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + "source": MagicMock(sdf=None), + }, + ] + builder = self._make_builder(bodies, shapes) + sdf_cfg = self._make_sdf_cfg() + + num_added, _ = NewtonManager._create_sdf_collision_from_visual(builder, {0}, sdf_cfg, None) + + assert num_added == 0 + builder.add_shape_mesh.assert_not_called() + + def test_warns_when_no_visual_mesh_source(self): + """A body with no mesh source logs a warning and is skipped.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/arm"] + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/arm/visual", + "geo_type": GeoType.MESH, + "flags": 0, + "source": None, + }, + ] + builder = self._make_builder(bodies, shapes) + sdf_cfg = self._make_sdf_cfg() + + with patch("isaaclab_newton.physics.newton_manager.logger") as mock_logger: + num_added, _ = NewtonManager._create_sdf_collision_from_visual(builder, {0}, sdf_cfg, None) + mock_logger.warning.assert_called() + + assert num_added == 0 + + def test_k_hydro_sets_hydroelastic_on_new_shape(self): + """When k_hydro is set, the new collision shape gets is_hydroelastic=True.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/arm"] + mesh = MagicMock(sdf=None) + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/arm/visual", + "geo_type": GeoType.MESH, + "flags": 0, + "source": mesh, + }, + ] + builder = self._make_builder(bodies, shapes) + sdf_cfg = self._make_sdf_cfg(k_hydro=1e10) + + num_added, num_hydro = NewtonManager._create_sdf_collision_from_visual(builder, {0}, sdf_cfg, None) + + assert num_added == 1 + assert num_hydro == 1 + call_kwargs = builder.add_shape_mesh.call_args[1] + assert call_kwargs["cfg"].is_hydroelastic is True + + def test_hydro_patterns_filter_respected(self): + """hydroelastic patterns filter which visual-mesh shapes get HYDROELASTIC.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + bodies = ["/World/Robot/arm", "/World/Robot/leg"] + mesh_arm = MagicMock(sdf=None) + mesh_leg = MagicMock(sdf=None) + shapes = [ + { + "body_idx": 0, + "label": "/World/Robot/arm/visual", + "geo_type": GeoType.MESH, + "flags": 0, + "source": mesh_arm, + }, + { + "body_idx": 1, + "label": "/World/Robot/leg/visual", + "geo_type": GeoType.MESH, + "flags": 0, + "source": mesh_leg, + }, + ] + builder = self._make_builder(bodies, shapes) + sdf_cfg = self._make_sdf_cfg(k_hydro=1e10) + hydro_patterns = [re.compile(".*arm.*")] + + num_added, num_hydro = NewtonManager._create_sdf_collision_from_visual( + builder, {0, 1}, sdf_cfg, None, hydro_patterns + ) + + assert num_added == 2 + assert num_hydro == 1 # only arm matches hydro pattern diff --git a/source/isaaclab_tasks/changelog.d/nut-thread-newton.minor.rst b/source/isaaclab_tasks/changelog.d/nut-thread-newton.minor.rst new file mode 100644 index 000000000000..a6bd3a538ac9 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/nut-thread-newton.minor.rst @@ -0,0 +1,15 @@ +Added +^^^^^ + +* Added Newton-backend support to :class:`~isaaclab_tasks.direct.factory.FactoryEnv`. + Run ``Isaac-Factory-*-Direct-v0`` tasks under the Newton physics backend with + ``presets=newton``; the existing PhysX path is unchanged. New helper modules: + + * :mod:`isaaclab_tasks.direct.factory.factory_control_newton` adapts + :class:`newton.selection.ArticulationView` to the J/M shapes the + Factory OSC math (``factory_control.compute_dof_torque``) expects, and + folds ``joint_armature`` into the mass-matrix diagonal. + * :mod:`isaaclab_tasks.direct.factory.factory_newton_setup` houses + procedural Newton-only post-load asset patches (POSITION-mode override, + parser-body gravity compensation, finger hydroelastic SDFs, contact-stack + tuning). diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_control.py b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_control.py index a43b7ffcc743..138d33ba1b50 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_control.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_control.py @@ -70,29 +70,30 @@ def compute_dof_torque( task_wrench.sign() * (task_wrench.abs() - dead_zone_thresholds), ) - # Set tau = J^T * tau, i.e., map tau into joint space as desired - jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2) - dof_torque[:, 0:7] = (jacobian_T @ task_wrench.unsqueeze(-1)).squeeze(-1) - # adapted from roboticsproceedings.org/rss07/p31.pdf # useful tensors - arm_mass_matrix_inv = torch.inverse(arm_mass_matrix) jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2) + arm_mass_matrix_inv = torch.inverse(arm_mass_matrix) arm_mass_matrix_task = torch.inverse( - jacobian @ torch.inverse(arm_mass_matrix) @ jacobian_T + jacobian @ arm_mass_matrix_inv @ jacobian_T ) # ETH eq. 3.86; geometric Jacobian is assumed - j_eef_inv = arm_mass_matrix_task @ jacobian @ arm_mass_matrix_inv - default_dof_pos_tensor = torch.tensor(cfg.ctrl.default_dof_pos_tensor, device=device).repeat((num_envs, 1)) - # nullspace computation - distance_to_default_dof_pos = default_dof_pos_tensor - dof_pos[:, :7] - distance_to_default_dof_pos = (distance_to_default_dof_pos + math.pi) % ( - 2 * math.pi - ) - math.pi # normalize to [-pi, pi] - u_null = cfg.ctrl.kd_null * -dof_vel[:, :7] + cfg.ctrl.kp_null * distance_to_default_dof_pos - u_null = arm_mass_matrix @ u_null.unsqueeze(-1) - torque_null = (torch.eye(7, device=device).unsqueeze(0) - torch.transpose(jacobian, 1, 2) @ j_eef_inv) @ u_null - dof_torque[:, 0:7] += torque_null.squeeze(-1) + + # Set tau = J^T * tau, i.e., map tau into joint space as desired + dof_torque[:, 0:7] = (jacobian_T @ task_wrench.unsqueeze(-1)).squeeze(-1) + + if not getattr(cfg.ctrl, "disable_nullspace", False): + j_eef_inv = arm_mass_matrix_task @ jacobian @ arm_mass_matrix_inv + default_dof_pos_tensor = torch.tensor(cfg.ctrl.default_dof_pos_tensor, device=device).repeat((num_envs, 1)) + # nullspace computation + distance_to_default_dof_pos = default_dof_pos_tensor - dof_pos[:, :7] + distance_to_default_dof_pos = (distance_to_default_dof_pos + math.pi) % ( + 2 * math.pi + ) - math.pi # normalize to [-pi, pi] + u_null = cfg.ctrl.kd_null * -dof_vel[:, :7] + cfg.ctrl.kp_null * distance_to_default_dof_pos + u_null = arm_mass_matrix @ u_null.unsqueeze(-1) + torque_null = (torch.eye(7, device=device).unsqueeze(0) - jacobian_T @ j_eef_inv) @ u_null + dof_torque[:, 0:7] += torque_null.squeeze(-1) # TODO: Verify it's okay to no longer do gripper control here. dof_torque = torch.clamp(dof_torque, min=-100.0, max=100.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_control_newton.py b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_control_newton.py new file mode 100644 index 000000000000..7e6bce042f3d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_control_newton.py @@ -0,0 +1,28 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Factory: per-joint effort clamp for the Newton backend. + +PR #5400's ``BaseArticulationData`` accessors cover OSC Jacobian + mass +matrix; only the effort clamp stays here because Newton's ``joint_f`` +write doesn't honour ``effort_limit_sim`` like PhysX's drive does. +""" + +from __future__ import annotations + +import torch + +# Franka FR3 per-joint torque limits [N·m] (datasheet, symmetric). +FRANKA_FR3_EFFORT_LIMITS: tuple[float, ...] = (87.0, 87.0, 87.0, 87.0, 12.0, 12.0, 12.0) + + +def clamp_to_effort_limits( + dof_torque: torch.Tensor, limits: tuple[float, ...] = FRANKA_FR3_EFFORT_LIMITS +) -> torch.Tensor: + """Clamp arm-DOF torques in-place; trailing DOFs (e.g. gripper) untouched.""" + n = len(limits) + lim = torch.as_tensor(limits, device=dof_torque.device, dtype=dof_torque.dtype) + dof_torque[..., :n] = torch.clamp(dof_torque[..., :n], min=-lim, max=lim) + return dof_torque diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py index c38fe071b161..e8434473effc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py @@ -5,11 +5,12 @@ import numpy as np import torch +import warp as wp import carb import isaaclab.sim as sim_utils -from isaaclab.assets import Articulation +from isaaclab.assets import Articulation, RigidObject, RigidObjectCfg from isaaclab.envs import DirectRLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane from isaaclab.utils import math as torch_utils @@ -19,6 +20,29 @@ from .factory_env_cfg import OBS_DIM_CFG, STATE_DIM_CFG, FactoryEnvCfg +def _set_sim_gravity(cfg: FactoryEnvCfg, gravity: tuple[float, float, float]) -> None: + """Set the live simulator gravity, dispatching on backend.""" + if _is_newton_backend(cfg): + from isaaclab_newton.physics import NewtonManager + from newton.solvers import SolverNotifyFlags + + if NewtonManager._model is not None: + NewtonManager._model.set_gravity(gravity) + # Re-upload model properties so the solver picks up the change. + NewtonManager.add_model_change(SolverNotifyFlags.MODEL_PROPERTIES) + return + physics_sim_view = sim_utils.SimulationContext.instance().physics_sim_view + physics_sim_view.set_gravity(carb.Float3(*gravity)) + + +def _is_newton_backend(cfg: FactoryEnvCfg) -> bool: + """Return True when the cfg has been resolved to the Newton physics backend.""" + physics = getattr(cfg.sim, "physics", None) + if physics is None: + return False + return type(physics).__module__.startswith("isaaclab_newton") + + class FactoryEnv(DirectRLEnv): cfg: FactoryEnvCfg @@ -29,13 +53,27 @@ def __init__(self, cfg: FactoryEnvCfg, render_mode: str | None = None, **kwargs) cfg.observation_space += cfg.action_space cfg.state_space += cfg.action_space self.cfg_task = cfg.task + self._is_newton = _is_newton_backend(cfg) + + if self._is_newton: + from . import factory_newton_setup + + factory_newton_setup.apply_cfg_overrides(cfg) super().__init__(cfg, render_mode, **kwargs) - factory_utils.set_body_inertias(self._robot, self.scene.num_envs) + # set_body_inertias / set_friction reach into PhysX root_view; not + # exposed on Newton's adapter — values from the spawn cfg take effect. + if not self._is_newton: + factory_utils.set_body_inertias(self._robot, self.scene.num_envs) self._init_tensors() self._set_default_dynamics_parameters() + if self._is_newton: + from . import factory_newton_setup + + factory_newton_setup.warm_up_kernels(self) + def _set_default_dynamics_parameters(self): """Set parameters defining dynamic interactions.""" self.default_gains = torch.tensor(self.cfg.ctrl.default_task_prop_gains, device=self.device).repeat( @@ -49,10 +87,11 @@ def _set_default_dynamics_parameters(self): (self.num_envs, 1) ) - # Set masses and frictions. - factory_utils.set_friction(self._held_asset, self.cfg_task.held_asset_cfg.friction, self.scene.num_envs) - factory_utils.set_friction(self._fixed_asset, self.cfg_task.fixed_asset_cfg.friction, self.scene.num_envs) - factory_utils.set_friction(self._robot, self.cfg_task.robot_cfg.friction, self.scene.num_envs) + # PhysX only — Newton's adapter doesn't expose root_view friction. + if not self._is_newton: + factory_utils.set_friction(self._held_asset, self.cfg_task.held_asset_cfg.friction, self.scene.num_envs) + factory_utils.set_friction(self._fixed_asset, self.cfg_task.fixed_asset_cfg.friction, self.scene.num_envs) + factory_utils.set_friction(self._robot, self.cfg_task.robot_cfg.friction, self.scene.num_envs) def _init_tensors(self): """Initialize tensors once.""" @@ -85,18 +124,38 @@ def _setup_scene(self): """Initialize simulation scene.""" spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg(), translation=(0.0, 0.0, -1.05)) - # spawn a usd file of a table into the scene - cfg = sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd") - cfg.func( - "/World/envs/env_.*/Table", cfg, translation=(0.55, 0.0, 0.0), orientation=(0.0, 0.0, 0.70711, 0.70711) - ) + # Newton: spawn a thin kinematic cuboid as the table top. + # The lab-table USD has no UsdPhysics.RigidBodyAPI, so it can't be + # a RigidObjectCfg directly. MeshCuboidCfg (not CuboidCfg) so + # _build_collision_sdfs can bake a voxel SDF. + if self._is_newton: + table_cfg = RigidObjectCfg( + prim_path="/World/envs/env_.*/Table", + spawn=sim_utils.MeshCuboidCfg( + size=(1.2, 0.6, 0.04), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.55, 0.0, -0.02), rot=(1.0, 0.0, 0.0, 0.0)), + ) + self._table = RigidObject(table_cfg) + else: + # PhysX path: keep the original Seattle-lab-table USD spawn. + cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd" + ) + cfg.func( + "/World/envs/env_.*/Table", cfg, translation=(0.55, 0.0, 0.0), orientation=(0.0, 0.0, 0.70711, 0.70711) + ) self._robot = Articulation(self.cfg.robot) - self._fixed_asset = Articulation(self.cfg_task.fixed_asset) - self._held_asset = Articulation(self.cfg_task.held_asset) + # ``class_type`` dispatch: PhysX → Articulation, Newton → RigidObject + # (after the Newton cfg conversion in ``__init__``). + self._fixed_asset = self.cfg_task.fixed_asset.class_type(self.cfg_task.fixed_asset) + self._held_asset = self.cfg_task.held_asset.class_type(self.cfg_task.held_asset) if self.cfg_task.name == "gear_mesh": - self._small_gear_asset = Articulation(self.cfg_task.small_gear_cfg) - self._large_gear_asset = Articulation(self.cfg_task.large_gear_cfg) + self._small_gear_asset = self.cfg_task.small_gear_cfg.class_type(self.cfg_task.small_gear_cfg) + self._large_gear_asset = self.cfg_task.large_gear_cfg.class_type(self.cfg_task.large_gear_cfg) self.scene.clone_environments(copy_from_source=False) if self.device == "cpu": @@ -104,16 +163,79 @@ def _setup_scene(self): self.scene.filter_collisions() self.scene.articulations["robot"] = self._robot - self.scene.articulations["fixed_asset"] = self._fixed_asset - self.scene.articulations["held_asset"] = self._held_asset + # rigid_objects on Newton, articulations on PhysX — mirrors shadow_hand_vision. + asset_registry = self.scene.rigid_objects if self._is_newton else self.scene.articulations + asset_registry["fixed_asset"] = self._fixed_asset + asset_registry["held_asset"] = self._held_asset if self.cfg_task.name == "gear_mesh": - self.scene.articulations["small_gear"] = self._small_gear_asset - self.scene.articulations["large_gear"] = self._large_gear_asset + asset_registry["small_gear"] = self._small_gear_asset + asset_registry["large_gear"] = self._large_gear_asset + + # Newton-only: register the kinematic table RigidObject for cloning. + if self._is_newton: + self.scene.rigid_objects["table"] = self._table # add lights light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) light_cfg.func("/World/Light", light_cfg) + # Register MODEL_INIT here so it fires before model finalization. + if self._is_newton: + from . import factory_newton_setup + + factory_newton_setup.register_model_init_callback() + + def _compute_fingertip_velocity_from_newton_state(self) -> tuple[torch.Tensor, torch.Tensor]: + """Read fingertip velocity from mjwarp state and transport COM → link origin. + + Newton's data adapter zeros ``body_lin_vel_w`` on the zero-mass + ``panda_fingertip_centered`` virtual link; without this helper + the OSC's task-space damping term goes to zero. + + Returns: + ``(linvel, angvel)`` each ``(num_envs, 3)`` in world frame at + the fingertip link origin. + """ + from isaaclab_newton.physics import NewtonManager + + state = NewtonManager._state_0 + model = NewtonManager._model + if state is None or model is None: + zero = torch.zeros((self.num_envs, 3), device=self.device) + return zero, zero + + # Cache the per-env fingertip body global indices on first call. + if not hasattr(self, "_fingertip_body_global_idx"): + labels = list(model.body_label) + idxs = [] + for env_idx in range(self.num_envs): + env_token = f"/env_{env_idx}/" + for i, lab in enumerate(labels): + if lab.endswith("/panda_fingertip_centered") and env_token in lab: + idxs.append(i) + break + if len(idxs) != self.num_envs: + raise RuntimeError( + f"Could not resolve panda_fingertip_centered for all envs (found {len(idxs)} of {self.num_envs})." + ) + self._fingertip_body_global_idx = torch.tensor(idxs, dtype=torch.long, device=self.device) + + # body_qd: [v_com_x, v_com_y, v_com_z, w_x, w_y, w_z], world frame. + body_qd_t = wp.to_torch(state.body_qd) + ft_idx = self._fingertip_body_global_idx + v_com = body_qd_t[ft_idx, 0:3] + omega = body_qd_t[ft_idx, 3:6] + + # v_link = v_com + omega × (p_link - p_com). + body_q_t = wp.to_torch(state.body_q) + body_com_t = wp.to_torch(model.body_com) + link_pos = body_q_t[ft_idx, 0:3] + link_quat_xyzw = body_q_t[ft_idx, 3:7] + com_world_offset = torch_utils.quat_apply(link_quat_xyzw, body_com_t[ft_idx]) + r_com_w = link_pos + com_world_offset + v_link = v_com + torch.cross(omega, link_pos - r_com_w, dim=-1) + return v_link, omega + def _compute_intermediate_values(self, dt): """Get values computed from raw tensors. This includes adding noise.""" # TODO: A lot of these can probably only be set once? @@ -127,11 +249,15 @@ def _compute_intermediate_values(self, dt): self._robot.data.body_pos_w.torch[:, self.fingertip_body_idx] - self.scene.env_origins ) self.fingertip_midpoint_quat = self._robot.data.body_quat_w.torch[:, self.fingertip_body_idx] - self.fingertip_midpoint_linvel = self._robot.data.body_lin_vel_w.torch[:, self.fingertip_body_idx] - self.fingertip_midpoint_angvel = self._robot.data.body_ang_vel_w.torch[:, self.fingertip_body_idx] + if self._is_newton: + self.fingertip_midpoint_linvel, self.fingertip_midpoint_angvel = ( + self._compute_fingertip_velocity_from_newton_state() + ) + else: + self.fingertip_midpoint_linvel = self._robot.data.body_lin_vel_w.torch[:, self.fingertip_body_idx] + self.fingertip_midpoint_angvel = self._robot.data.body_ang_vel_w.torch[:, self.fingertip_body_idx] jacobians = self._robot.data.body_link_jacobian_w.torch - self.left_finger_jacobian = jacobians[:, self.left_finger_body_idx - 1, 0:6, 0:7] self.right_finger_jacobian = jacobians[:, self.right_finger_body_idx - 1, 0:6, 0:7] self.fingertip_midpoint_jacobian = (self.left_finger_jacobian + self.right_finger_jacobian) * 0.5 @@ -198,6 +324,21 @@ def _get_observations(self): obs_tensors = factory_utils.collapse_obs_dict(obs_dict, self.cfg.obs_order + ["prev_actions"]) state_tensors = factory_utils.collapse_obs_dict(state_dict, self.cfg.state_order + ["prev_actions"]) + + # On Newton, NaN-divergent worlds occasionally surface in obs/state + # tensors and poison the policy's normal distribution head ("std must + # be >= 0"). Detect, substitute zeros so rl_games keeps running, and + # delegate solver-side scrub to NewtonManager (which also queues the + # affected envs for a second-pass sanitize at the next reset). + if self._is_newton: + from isaaclab_newton.physics import NewtonManager # noqa: PLC0415 + + nan_mask = torch.isnan(obs_tensors).any(dim=-1) | torch.isnan(state_tensors).any(dim=-1) + if nan_mask.any(): + NewtonManager.flag_nan_envs(nan_mask) + obs_tensors = torch.nan_to_num(obs_tensors, nan=0.0) + state_tensors = torch.nan_to_num(state_tensors, nan=0.0) + return {"policy": obs_tensors, "critic": state_tensors} def _reset_buffers(self, env_ids): @@ -328,6 +469,12 @@ def generate_ctrl_signals( self.ctrl_target_joint_pos[:, 7:9] = ctrl_target_gripper_dof_pos self.joint_torque[:, 7:9] = 0.0 + # Newton's joint_f writes don't honour effort_limit_sim — clamp explicitly. + if self._is_newton: + from . import factory_control_newton + + factory_control_newton.clamp_to_effort_limits(self.joint_torque) + self._robot.set_joint_position_target_index(target=self.ctrl_target_joint_pos) self._robot.set_joint_effort_target_index(target=self.joint_torque) @@ -419,6 +566,12 @@ def _get_rewards(self): for rew_name, rew in rew_dict.items(): rew_buf += rew_dict[rew_name] * rew_scales[rew_name] + # On Newton, flagged NaN-divergent worlds may produce NaN rewards + # between detection (in _get_observations) and the next reset. + # Substitute zeros so the rolling reward signal isn't poisoned. + if self._is_newton: + rew_buf = torch.nan_to_num(rew_buf, nan=0.0) + self.prev_actions = self.actions.clone() self._log_factory_metrics(rew_dict, curr_successes) @@ -490,6 +643,11 @@ def _get_factory_rew_dict(self, curr_successes): def _reset_idx(self, env_ids): """We assume all envs will always be reset at the same time.""" + if self._is_newton: + from isaaclab_newton.physics import NewtonManager # noqa: PLC0415 + + NewtonManager.sanitize_pending_nan_envs() + super()._reset_idx(env_ids) self._set_assets_to_default_pose(env_ids) @@ -504,16 +662,16 @@ def _set_assets_to_default_pose(self, env_ids): held_vel = self._held_asset.data.default_root_vel.torch.clone()[env_ids] held_pose[:, 0:3] += self.scene.env_origins[env_ids] held_vel[:] = 0.0 - self._held_asset.write_root_pose_to_sim_index(root_pose=held_pose, env_ids=env_ids) - self._held_asset.write_root_velocity_to_sim_index(root_velocity=held_vel, env_ids=env_ids) + self._held_asset.write_root_link_pose_to_sim_index(root_pose=held_pose, env_ids=env_ids) + self._held_asset.write_root_link_velocity_to_sim_index(root_velocity=held_vel, env_ids=env_ids) self._held_asset.reset() fixed_pose = self._fixed_asset.data.default_root_pose.torch.clone()[env_ids] fixed_vel = self._fixed_asset.data.default_root_vel.torch.clone()[env_ids] fixed_pose[:, 0:3] += self.scene.env_origins[env_ids] fixed_vel[:] = 0.0 - self._fixed_asset.write_root_pose_to_sim_index(root_pose=fixed_pose, env_ids=env_ids) - self._fixed_asset.write_root_velocity_to_sim_index(root_velocity=fixed_vel, env_ids=env_ids) + self._fixed_asset.write_root_link_pose_to_sim_index(root_pose=fixed_pose, env_ids=env_ids) + self._fixed_asset.write_root_link_velocity_to_sim_index(root_velocity=fixed_vel, env_ids=env_ids) self._fixed_asset.reset() def set_pos_inverse_kinematics( @@ -550,7 +708,7 @@ def set_pos_inverse_kinematics( self._robot.write_joint_velocity_to_sim_index(velocity=self.joint_vel) self._robot.set_joint_position_target_index(target=self.ctrl_target_joint_pos) - # Simulate and update tensors. + # PhysX steps physics; Newton refreshes FK + Jacobian only. self.step_sim_no_action() ik_time += self.physics_dt @@ -612,6 +770,16 @@ def step_sim_no_action(self): This method should only be called during resets when all environments reset at the same time. + + Both backends now run the full ``sim.step``. An earlier optimization + attempt skipped the integrator on Newton in favor of a direct + ``eval_fk`` refresh, but that bypassed Newton's captured-CUDA-graph + scatter/gather between staging arrays and ``state_0``, leaving the + IsaacLab data-layer bindings (``body_pos_w`` etc.) reading stale + values. DLS IK then diverged because the Jacobian and the cached + fingertip pose disagreed about the current state. The captured + graph is fast (~1.5 ms after warm-up), so paying it per IK + iteration is acceptable. """ self.scene.write_data_to_sim() self.sim.step(render=False) @@ -620,9 +788,13 @@ def step_sim_no_action(self): def randomize_initial_state(self, env_ids): """Randomize initial state and perform any episode-level randomization.""" - # Disable gravity. - physics_sim_view = sim_utils.SimulationContext.instance().physics_sim_view - physics_sim_view.set_gravity(carb.Float3(0.0, 0.0, 0.0)) + self._full_reset(env_ids) + + def _full_reset(self, env_ids): + """Original PhysX reset path: IK + asset randomization + grasp settle.""" + + # Disable gravity (PhysX/Newton-portable). + _set_sim_gravity(self.cfg, (0.0, 0.0, 0.0)) # (1.) Randomize fixed asset pose. fixed_pose = self._fixed_asset.data.default_root_pose.torch.clone()[env_ids] @@ -648,8 +820,8 @@ def randomize_initial_state(self, env_ids): # (1.c.) Velocity fixed_vel[:] = 0.0 # vel # (1.d.) Update values. - self._fixed_asset.write_root_pose_to_sim_index(root_pose=fixed_pose, env_ids=env_ids) - self._fixed_asset.write_root_velocity_to_sim_index(root_velocity=fixed_vel, env_ids=env_ids) + self._fixed_asset.write_root_link_pose_to_sim_index(root_pose=fixed_pose, env_ids=env_ids) + self._fixed_asset.write_root_link_velocity_to_sim_index(root_velocity=fixed_vel, env_ids=env_ids) self._fixed_asset.reset() # (1.e.) Noisy position observation. @@ -737,16 +909,16 @@ def randomize_initial_state(self, env_ids): small_gear_vel = self._small_gear_asset.data.default_root_vel.torch.clone()[env_ids] small_gear_pose[:, 0:7] = fixed_pose[:, 0:7] small_gear_vel[:] = 0.0 # vel - self._small_gear_asset.write_root_pose_to_sim_index(root_pose=small_gear_pose, env_ids=env_ids) - self._small_gear_asset.write_root_velocity_to_sim_index(root_velocity=small_gear_vel, env_ids=env_ids) + self._small_gear_asset.write_root_link_pose_to_sim_index(root_pose=small_gear_pose, env_ids=env_ids) + self._small_gear_asset.write_root_link_velocity_to_sim_index(root_velocity=small_gear_vel, env_ids=env_ids) self._small_gear_asset.reset() large_gear_pose = self._large_gear_asset.data.default_root_pose.torch.clone()[env_ids] large_gear_vel = self._large_gear_asset.data.default_root_vel.torch.clone()[env_ids] large_gear_pose[:, 0:7] = fixed_pose[:, 0:7] large_gear_vel[:] = 0.0 # vel - self._large_gear_asset.write_root_pose_to_sim_index(root_pose=large_gear_pose, env_ids=env_ids) - self._large_gear_asset.write_root_velocity_to_sim_index(root_velocity=large_gear_vel, env_ids=env_ids) + self._large_gear_asset.write_root_link_pose_to_sim_index(root_pose=large_gear_pose, env_ids=env_ids) + self._large_gear_asset.write_root_link_velocity_to_sim_index(root_velocity=large_gear_vel, env_ids=env_ids) self._large_gear_asset.reset() # (3) Randomize asset-in-gripper location. @@ -789,8 +961,8 @@ def randomize_initial_state(self, env_ids): held_pose[:, 0:3] = translated_held_asset_pos + self.scene.env_origins held_pose[:, 3:7] = translated_held_asset_quat held_vel[:] = 0.0 - self._held_asset.write_root_pose_to_sim_index(root_pose=held_pose) - self._held_asset.write_root_velocity_to_sim_index(root_velocity=held_vel) + self._held_asset.write_root_link_pose_to_sim_index(root_pose=held_pose) + self._held_asset.write_root_link_velocity_to_sim_index(root_velocity=held_vel) self._held_asset.reset() # Close hand @@ -828,4 +1000,8 @@ def randomize_initial_state(self, env_ids): self.task_prop_gains = self.default_gains self.task_deriv_gains = factory_utils.get_deriv_gains(self.default_gains) - physics_sim_view.set_gravity(carb.Float3(*self.cfg.sim.gravity)) + # Newton ignores per-body disable_gravity flags — keep global gravity at zero. + if self._is_newton: + _set_sim_gravity(self.cfg, (0.0, 0.0, 0.0)) + else: + _set_sim_gravity(self.cfg, tuple(self.cfg.sim.gravity)) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env_cfg.py index a250380e2566..cfe32f643ee0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env_cfg.py @@ -3,6 +3,12 @@ # # SPDX-License-Identifier: BSD-3-Clause +from isaaclab_newton.physics import ( + HydroelasticSDFCfg, + MJWarpSolverCfg, + NewtonCfg, + NewtonCollisionPipelineCfg, +) from isaaclab_physx.physics import PhysxCfg import isaaclab.sim as sim_utils @@ -14,6 +20,8 @@ from isaaclab.sim.spawners.materials.physics_materials_cfg import RigidBodyMaterialCfg from isaaclab.utils.configclass import configclass +from isaaclab_tasks.utils import PresetCfg + from .factory_tasks_cfg import ASSET_DIR, FactoryTask, GearMesh, NutThread, PegInsert OBS_DIM_CFG = { @@ -68,6 +76,98 @@ class CtrlCfg: kp_null = 10.0 kd_null = 6.3246 + # When True, skip the OSC null-space term in :func:`factory_control.compute_dof_torque`. + # Default ``False`` keeps the original Factory PhysX behaviour (null-space on); + # :meth:`FactoryEnv.__init__` overrides it to ``True`` on Newton because the + # null-space target ``default_dof_pos_tensor`` does not match the IK-converged + # "above bolt" arm config the env actually holds, so leaving null-space on + # leaks a constant TCP-direction torque under Newton's mjwarp OSC integration. + disable_nullspace: bool = False + + +_NEWTON_SOLVER_CFG = MJWarpSolverCfg( + solver="newton", + integrator="implicitfast", + njmax=4000, + nconmax=4000, + # Higher impratio over-constrains the gripper pad on the nut. + impratio=10.0, + cone="elliptic", + use_mujoco_contacts=False, + iterations=10, + ls_iterations=100, +) + + +@configclass +class FactoryPhysicsCfg(PresetCfg): + """Per-backend physics cfg. Newton offers two SDF modes. + + - ``newton``: SDF mesh contacts with hydroelastic distributed pressure. + ``num_substeps=8``, ``collision_decimation=2`` (re-collide 4x per tick). + - ``newton_sdf``: SDF mesh contacts with vanilla penalty-spring forces + (no hydroelastic). The stiffer normal response requires denser + substepping — ``num_substeps=10``, ``collision_decimation=1`` + (re-collide every substep) — to keep contact normals fresh. + """ + + physx = PhysxCfg( + solver_type=1, + max_position_iteration_count=192, # Important to avoid interpenetration. + max_velocity_iteration_count=1, + bounce_threshold_velocity=0.2, + friction_offset_threshold=0.01, + friction_correlation_distance=0.00625, + gpu_max_rigid_contact_count=2**23, + gpu_max_rigid_patch_count=2**23, + gpu_collision_stack_size=2**28, + gpu_max_num_partitions=1, # Important for stable simulation. + ) + newton = NewtonCfg( + solver_cfg=_NEWTON_SOLVER_CFG, + collision_cfg=NewtonCollisionPipelineCfg( + broad_phase="explicit", + rigid_contact_max=32768, + # PhysX patch-friction analogs. + sdf_hydroelastic_config=HydroelasticSDFCfg( + anchor_contact=True, + moment_matching=True, + output_contact_surface=False, + ), + ), + # 1.04 ms substep dt; re-collide every 2 substeps (4x per tick). + num_substeps=8, + collision_decimation=2, + use_cuda_graph=True, + ) + newton_sdf = NewtonCfg( + solver_cfg=_NEWTON_SOLVER_CFG, + collision_cfg=NewtonCollisionPipelineCfg( + broad_phase="explicit", + rigid_contact_max=32768, + # No hydroelastic — fall back to penalty-spring contacts. + sdf_hydroelastic_config=None, + ), + # Vanilla SDF needs more substeps than hydroelastic. + num_substeps=10, + collision_decimation=1, + use_cuda_graph=True, + ) + default = physx + + +@configclass +class FactorySceneCfg(PresetCfg): + """Per-backend scene cfg. + + PhysX uses Fabric-layer scene cloning for throughput; Newton does not + support ``clone_in_fabric=True`` and must use the per-prim clone path. + """ + + physx: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=128, env_spacing=2.0, clone_in_fabric=True) + newton: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=128, env_spacing=2.0, clone_in_fabric=False) + default: InteractiveSceneCfg = physx + @configclass class FactoryEnvCfg(DirectRLEnvCfg): @@ -100,25 +200,14 @@ class FactoryEnvCfg(DirectRLEnvCfg): device="cuda:0", dt=1 / 120, gravity=(0.0, 0.0, -9.81), - physics=PhysxCfg( - solver_type=1, - max_position_iteration_count=192, # Important to avoid interpenetration. - max_velocity_iteration_count=1, - bounce_threshold_velocity=0.2, - friction_offset_threshold=0.01, - friction_correlation_distance=0.00625, - gpu_max_rigid_contact_count=2**23, - gpu_max_rigid_patch_count=2**23, - gpu_collision_stack_size=2**28, - gpu_max_num_partitions=1, # Important for stable simulation. - ), + physics=FactoryPhysicsCfg(), physics_material=RigidBodyMaterialCfg( static_friction=1.0, dynamic_friction=1.0, ), ) - scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=128, env_spacing=2.0, clone_in_fabric=True) + scene: InteractiveSceneCfg = FactorySceneCfg() robot = ArticulationCfg( prim_path="/World/envs/env_.*/Robot", diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_newton_setup.py b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_newton_setup.py new file mode 100644 index 000000000000..fa6d690f0e1a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_newton_setup.py @@ -0,0 +1,454 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Factory Newton-only setup: cfg overrides, model-init callback, kernel warm-up.""" + +from __future__ import annotations + +import logging +import os +from typing import TYPE_CHECKING + +import newton + +from isaaclab.assets import ArticulationCfg, RigidObjectCfg + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from .factory_env import FactoryEnv + from .factory_env_cfg import FactoryEnvCfg + +_ARM_JOINT_NAMES = [f"panda_joint{i}" for i in range(1, 8)] +_FINGER_JOINT_NAMES = ["panda_finger_joint1", "panda_finger_joint2"] +_ROBOT_BODY_PATH_SUBSTR = "/Robot/" + +# Latched in :func:`apply_cfg_overrides`; read by the MODEL_INIT callback. +_use_hydroelastic: bool = True + + +def register_model_init_callback() -> None: + """Wire a Newton MODEL_INIT callback that mutates the builder.""" + from isaaclab_newton.physics import NewtonManager + + from isaaclab.physics import PhysicsEvent + + NewtonManager.register_callback( + lambda _ev: _model_init_callback(), PhysicsEvent.MODEL_INIT, name="factory_newton_setup" + ) + + +def apply_cfg_overrides(cfg: FactoryEnvCfg) -> None: + """Apply Newton-only cfg overrides in place, before ``super().__init__``.""" + cfg.sim.dt = 1.0 / 120.0 + cfg.decimation = 8 + cfg.ctrl.disable_nullspace = True + + cfg_task = cfg.task + kinematic_for = {"fixed_asset", "small_gear_cfg", "large_gear_cfg"} + for asset_attr in ("fixed_asset", "held_asset", "small_gear_cfg", "large_gear_cfg"): + asset_cfg = getattr(cfg_task, asset_attr, None) + if isinstance(asset_cfg, ArticulationCfg): + setattr(cfg_task, asset_attr, _to_rigid_object_cfg(asset_cfg, kinematic=asset_attr in kinematic_for)) + + # Zero the bolt reset's Z noise: the bolt is a kinematic body (USD + # PhysicsFixedJoint) so it sticks wherever the reset places it. + noise = getattr(cfg_task, "fixed_asset_init_pos_noise", None) + if noise is not None and len(noise) >= 3: + cfg_task.fixed_asset_init_pos_noise = list(noise[:2]) + [0.0] + + # Bolt init z=0.0 matches the Newton cuboid table top. + bolt = getattr(cfg_task, "fixed_asset", None) + if bolt is not None and hasattr(bolt, "init_state") and hasattr(bolt.init_state, "pos"): + cur_pos = bolt.init_state.pos + new_init = bolt.init_state.replace(pos=(float(cur_pos[0]), float(cur_pos[1]), 0.0)) + cfg_task.fixed_asset = bolt.replace(init_state=new_init) + + # Arm armature stabilises OSC Lambda — must be set pre-super().__init__. + arm_actuators = cfg.robot.actuators + arm_actuators["panda_arm1"].armature = 0.3 + arm_actuators["panda_arm2"].armature = 0.11 + arm_actuators["panda_hand"].armature = 0.15 + + # OSC kd already damps via 2√Kp + armature; extra joint kd brakes tracking. + arm_actuators["panda_arm1"].damping = 0.0 + arm_actuators["panda_arm2"].damping = 0.0 + + global _use_hydroelastic + _use_hydroelastic = getattr(cfg.sim.physics.collision_cfg, "sdf_hydroelastic_config", None) is not None + + if _use_hydroelastic: + arm_actuators["panda_hand"].stiffness = 1000.0 + arm_actuators["panda_hand"].damping = 10.0 + else: + kp = float(os.environ.get("FACTORY_SDF_FINGER_KP", "8000")) + kd = float(os.environ.get("FACTORY_SDF_FINGER_KD", "160")) + arm_actuators["panda_hand"].stiffness = kp + arm_actuators["panda_hand"].damping = kd + + solver_cfg = getattr(cfg.sim.physics, "solver_cfg", None) + if solver_cfg is not None: + iters_override = int(os.environ.get("FACTORY_SOLVER_ITERATIONS", "0")) + imp_override = float(os.environ.get("FACTORY_SOLVER_IMPRATIO", "0")) + if iters_override > 0 and hasattr(solver_cfg, "iterations"): + solver_cfg.iterations = iters_override + if imp_override > 0 and hasattr(solver_cfg, "impratio"): + solver_cfg.impratio = imp_override + + substeps_override = int(os.environ.get("FACTORY_NUM_SUBSTEPS", "0")) + if substeps_override > 0 and hasattr(cfg.sim.physics, "num_substeps"): + cfg.sim.physics.num_substeps = substeps_override + + _monkey_patch_cloner_no_simplify() + + # Scale contact buffer with num_envs; SDF-only emits ~3x hydroelastic. + if ( + getattr(cfg.sim.physics, "collision_cfg", None) is not None + and getattr(cfg.sim.physics, "solver_cfg", None) is not None + and hasattr(cfg.sim.physics.solver_cfg, "njmax") + ): + njmax_mult = float(os.environ.get("FACTORY_NJMAX_MULT", "1.0")) + njmax = int(cfg.sim.physics.solver_cfg.njmax) + num_worlds = int(cfg.scene.num_envs) + cfg.sim.physics.collision_cfg.rigid_contact_max = int(njmax * num_worlds * njmax_mult) + + +def warm_up_kernels(env: FactoryEnv) -> None: + """Pre-JIT Newton's per-step kernels so the reset IK loop runs on warm caches.""" + for _ in range(2): + env.step_sim_no_action() + + +def _to_rigid_object_cfg(art_cfg: ArticulationCfg, kinematic: bool) -> RigidObjectCfg: + """Adapt a joint-less Factory ArticulationCfg into a RigidObjectCfg.""" + import isaaclab.sim as sim_utils # noqa: PLC0415 + + rigid_props = art_cfg.spawn.rigid_props + if kinematic: + rigid_props = ( + rigid_props.replace(kinematic_enabled=True) + if rigid_props is not None + else sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True) + ) + spawn = art_cfg.spawn.replace(rigid_props=rigid_props) + return RigidObjectCfg( + prim_path=art_cfg.prim_path, + spawn=spawn, + init_state=RigidObjectCfg.InitialStateCfg( + pos=art_cfg.init_state.pos, + rot=art_cfg.init_state.rot, + lin_vel=art_cfg.init_state.lin_vel, + ang_vel=art_cfg.init_state.ang_vel, + ), + ) + + +def _monkey_patch_cloner_no_simplify() -> None: + """Skip the cloner's convex-hull mesh approximation so per-shape SDFs survive.""" + from isaaclab_newton.cloner import newton_replicate as _nr + + _orig = _nr.newton_physics_replicate + + def _no_simplify(*args, **kwargs): + kwargs.setdefault("simplify_meshes", False) + return _orig(*args, **kwargs) + + _nr.newton_physics_replicate = _no_simplify + import isaaclab_newton.cloner as _cloner_module # noqa: PLC0415 + + _cloner_module.newton_physics_replicate = _no_simplify + + +# --------------------------------------------------------------------------- +# Builder-time setup (MODEL_INIT callback body). +# --------------------------------------------------------------------------- + + +def _model_init_callback() -> None: + """Body of the MODEL_INIT callback. Operates on the live builder.""" + from isaaclab_newton.physics import NewtonManager + + builder = NewtonManager._builder + if builder is None: + return + + _set_joint_target_mode(builder) + _set_ctrl_source_joint_target(builder) + _filter_base_table_contacts(builder) + _tune_nut_bolt_contacts(builder) + _build_collision_sdfs(builder) + # Armature is applied via ImplicitActuatorCfg, not the builder — finalize + # overwrites builder writes at robot DOF indices with the actuator value. + + +def _joint_label_indices(builder, name_substrs: list[str]) -> list[int]: + """Return joint indices whose label contains any of ``name_substrs``.""" + return [i for i, label in enumerate(builder.joint_label) if any(s in label for s in name_substrs)] + + +def _joint_indices_to_dof_indices(builder, joint_idxs: list[int]) -> list[int]: + """Translate joint indices into DOF indices via ``joint_qd_start``. + + Free-floating joints contribute 6 qd entries per joint, so + ``joint_index != dof_index`` past env 0; walk ``joint_qd_start`` to + cover every DOF the joint owns. + """ + qd_start = list(builder.joint_qd_start) + total_dofs = qd_start[-1] if qd_start else 0 + out: list[int] = [] + for j in joint_idxs: + dof_lo = int(qd_start[j]) + dof_hi = int(qd_start[j + 1]) if (j + 1) < len(qd_start) else int(total_dofs) + out.extend(range(dof_lo, dof_hi)) + return out + + +def _arm_dof_indices(builder) -> list[int]: + """DOF indices in ``builder.joint_target_mode`` for the arm joints.""" + js = _joint_label_indices(builder, _ARM_JOINT_NAMES) + return _joint_indices_to_dof_indices(builder, js) + + +def _finger_dof_indices(builder) -> list[int]: + """DOF indices in ``builder.joint_target_mode`` for the finger joints.""" + js = _joint_label_indices(builder, _FINGER_JOINT_NAMES) + return _joint_indices_to_dof_indices(builder, js) + + +def _robot_body_indices(builder) -> list[int]: + """Indices into ``builder.body_label`` that belong to the Franka.""" + return [i for i, label in enumerate(builder.body_label) if _ROBOT_BODY_PATH_SUBSTR in label] + + +def _set_joint_target_mode(builder) -> None: + """Force every robot DOF (arm + fingers) into POSITION mode.""" + for dof_idx in _arm_dof_indices(builder) + _finger_dof_indices(builder): + builder.joint_target_mode[dof_idx] = int(newton.JointTargetMode.POSITION) + + +def _set_ctrl_source_joint_target(builder) -> None: + """Pin every robot DOF's ``mujoco:ctrl_source`` to ``JOINT_TARGET``. + + Required so mjwarp reads targets from ``Control.joint_target_pos`` + (set by ``set_joint_position_target_index``) rather than the unused + ``Control.mujoco.ctrl`` array. + """ + from newton.solvers import SolverMuJoCo + + custom = builder.custom_attributes.get("mujoco:ctrl_source") + if custom is None: + return + n_acts = len(builder.joint_target_mode) + if custom.values is None or len(custom.values) != n_acts: + custom.values = [int(SolverMuJoCo.CtrlSource.JOINT_TARGET)] * n_acts + else: + target_value = int(SolverMuJoCo.CtrlSource.JOINT_TARGET) + for dof_idx in _arm_dof_indices(builder) + _finger_dof_indices(builder): + custom.values[dof_idx] = target_value + + +def _filter_base_table_contacts(builder) -> None: + """Filter spurious robot-base ↔ table contacts. + + Without this filter mjwarp generates continuous tiny contacts between + base links and the table top; they propagate through the chain and + show up as multi-millimeter TCP drift during OSC hold. + """ + base_suffixes = ("/panda_link0", "/panda_link1") + table_substr = "/Table" + base_shape_idxs = [ + i + for i, body_idx in enumerate(builder.shape_body) + if 0 <= body_idx < len(builder.body_label) and builder.body_label[body_idx].endswith(base_suffixes) + ] + table_shape_idxs = [ + i + for i, body_idx in enumerate(builder.shape_body) + if 0 <= body_idx < len(builder.body_label) and table_substr in builder.body_label[body_idx] + ] + if not base_shape_idxs or not table_shape_idxs: + return + for base_i in base_shape_idxs: + for table_i in table_shape_idxs: + pair = (min(base_i, table_i), max(base_i, table_i)) + builder.shape_collision_filter_pairs.append(pair) + logger.info("Filtered %d base<->table collision pairs.", len(base_shape_idxs) * len(table_shape_idxs)) + + +def _tune_nut_bolt_contacts(builder) -> None: + """Tune contact-material gains on every nut/bolt collision shape. + + Hydroelastic mode keeps the penalty spring soft so it doesn't + double-count the hydroelastic normal force; SDF-only inherits the + Newton nut/bolt example values (ke=1e7, kd=1e4, gap=5 mm). + """ + if not all(hasattr(builder, attr) for attr in ("shape_label", "shape_material_mu", "shape_gap")): + return + + if _use_hydroelastic: + ke, kd, gap = 1.0e4, 100.0, 0.0 + else: + ke = float(os.environ.get("FACTORY_SDF_NUT_BOLT_KE", "1.0e7")) + kd = float(os.environ.get("FACTORY_SDF_NUT_BOLT_KD", "1.0e4")) + gap = float(os.environ.get("FACTORY_SDF_NUT_BOLT_GAP", "0.005")) + + for i in range(builder.shape_count): + label = str(builder.shape_label[i]).lower() + if "nut" in label: + builder.shape_material_mu[i] = 0.2 + builder.shape_material_ke[i] = ke + builder.shape_material_kd[i] = kd + builder.shape_gap[i] = gap + elif "bolt" in label: + builder.shape_material_mu[i] = 0.5 + builder.shape_material_ke[i] = ke + builder.shape_material_kd[i] = kd + builder.shape_gap[i] = gap + + +# --------------------------------------------------------------------------- +# SDF collision setup. Both modes build the same SDF grids; only the +# per-shape material flags differ: +# * hydroelastic: HYDROELASTIC flag + ``kh`` on finger/nut/bolt shapes. +# * sdf-only: ``ke`` / ``kd`` penalty springs on finger shapes. +# Finger friction extras (torsional + condim=4) apply in both modes. +# --------------------------------------------------------------------------- + +_SDF_RES_FINGER = 192 +_SDF_RES_NUT_BOLT = 256 +_SDF_RES_PANDA = 64 +# Table is large + flat (1.2 × 0.6 × 0.04 m). 32³ → ~37 mm cells in xy, +# ~1.25 mm in z — plenty for a "rest on a flat surface" SDF without +# burning memory on a high-res grid. +_SDF_RES_TABLE = 32 +_SDF_BAND_FINGER = (-0.01, 0.01) +_SDF_BAND_NUT_BOLT = (-0.005, 0.005) +_SDF_BAND_PANDA = (-0.01, 0.01) +# Wider outside band so bolt + nut see the table SDF from slightly above. +_SDF_BAND_TABLE = (-0.005, 0.02) + +# Hydroelastic stiffness [Pa/m]. Smaller values let the finger SDF punch +# through the nut SDF under the ~187 N PD close force. +_KH_FINGER = 1e11 +_KH_NUT_BOLT = 1e11 + +# SDF-only finger penalty spring. The Newton nut/bolt example ships +# ke=1e7/kd=1e4, but its scene has no gripper; finger pads under the +# saturated 40 N close force need ~30× stiffer contact. Sweep picked +# (3e8, 3e5) on Factory_develop_baseline.pth (success 0.625 → 0.891); +# higher values oscillate. +_KE_FINGER = 3.0e8 +_KD_FINGER = 3.0e5 + +_FINGER_MU_TORSIONAL = 0.1 +_FINGER_CONDIM = 4 + + +def _build_collision_sdfs(builder) -> None: + """Bake per-shape SDFs and tag finger/nut/bolt for the active contact model. + + Idempotent: meshes with a populated ``sdf`` are skipped, so re-runs + on a hot builder don't rebake. + """ + finger_names = ("panda_leftfinger", "panda_rightfinger") + finger_body_idxs = {i for i, label in enumerate(builder.body_label) if any(n in label for n in finger_names)} + nut_body_idxs = {i for i, label in enumerate(builder.body_label) if "HeldAsset/factory_nut_loose" in label} + bolt_body_idxs = {i for i, label in enumerate(builder.body_label) if "FixedAsset/factory_bolt_loose" in label} + table_body_idxs = {i for i, label in enumerate(builder.body_label) if "/Table" in label} + panda_body_idxs = set(_robot_body_indices(builder)) - finger_body_idxs + + meshlike = (newton.GeoType.MESH, newton.GeoType.CONVEX_MESH) + counts = { + "finger": 0, + "nut": 0, + "bolt": 0, + "panda": 0, + "table": 0, + "skip_no_mesh": 0, + "skip_already_built": 0, + } + + condim_attr = builder.custom_attributes.get("mujoco:condim") + if condim_attr is not None and condim_attr.values is None: + condim_attr.values = {} + + for shape_idx, body_idx in enumerate(builder.shape_body): + if int(builder.shape_type[shape_idx]) not in (int(t) for t in meshlike): + continue + if not (int(builder.shape_flags[shape_idx]) & int(newton.ShapeFlags.COLLIDE_SHAPES)): + continue + + mesh = builder.shape_source[shape_idx] + if mesh is None: + counts["skip_no_mesh"] += 1 + continue + + if body_idx in finger_body_idxs: + category = "finger" + res, band = _SDF_RES_FINGER, _SDF_BAND_FINGER + elif body_idx in nut_body_idxs: + category = "nut" + res, band = _SDF_RES_NUT_BOLT, _SDF_BAND_NUT_BOLT + elif body_idx in bolt_body_idxs: + category = "bolt" + res, band = _SDF_RES_NUT_BOLT, _SDF_BAND_NUT_BOLT + elif body_idx in panda_body_idxs: + category = "panda" + res, band = _SDF_RES_PANDA, _SDF_BAND_PANDA + elif body_idx in table_body_idxs: + # Voxel SDF for "Show Collision"; not HYDROELASTIC — keeps + # bolt-on-table contact off the hydroelastic mass balance. + category = "table" + res, band = _SDF_RES_TABLE, _SDF_BAND_TABLE + else: + continue + + if mesh.sdf is None: + # build_sdf doesn't honour shape_scale, so bake it into vertices. + shape_scale = builder.shape_scale[shape_idx] + scale_arr = (float(shape_scale[0]), float(shape_scale[1]), float(shape_scale[2])) + if not ( + abs(scale_arr[0] - 1.0) < 1e-6 and abs(scale_arr[1] - 1.0) < 1e-6 and abs(scale_arr[2] - 1.0) < 1e-6 + ): + import numpy as _np # noqa: PLC0415 + + scaled_verts = mesh.vertices * _np.asarray(scale_arr, dtype=_np.float32) + mesh = mesh.copy(vertices=scaled_verts, recompute_inertia=True) + builder.shape_source[shape_idx] = mesh + builder.shape_scale[shape_idx] = (1.0, 1.0, 1.0) + mesh.build_sdf(max_resolution=res, narrow_band_range=band, margin=abs(band[1])) + counts[category] += 1 + else: + counts["skip_already_built"] += 1 + + if _use_hydroelastic and category in ("finger", "nut", "bolt"): + builder.shape_flags[shape_idx] |= int(newton.ShapeFlags.HYDROELASTIC) + builder.shape_material_kh[shape_idx] = _KH_FINGER if category == "finger" else _KH_NUT_BOLT + elif not _use_hydroelastic and category == "finger": + builder.shape_material_ke[shape_idx] = float( + os.environ.get("FACTORY_SDF_FINGER_CONTACT_KE", str(_KE_FINGER)) + ) + builder.shape_material_kd[shape_idx] = float( + os.environ.get("FACTORY_SDF_FINGER_CONTACT_KD", str(_KD_FINGER)) + ) + finger_gap_env = os.environ.get("FACTORY_SDF_FINGER_GAP") + if finger_gap_env is not None: + builder.shape_gap[shape_idx] = float(finger_gap_env) + if category == "finger": + builder.shape_material_mu_torsional[shape_idx] = _FINGER_MU_TORSIONAL + if condim_attr is not None: + condim_attr.values[shape_idx] = _FINGER_CONDIM + + logger.info( + "Built SDFs (%s): finger=%d nut=%d bolt=%d panda=%d table=%d (skipped: no_mesh=%d, already_built=%d).", + "hydroelastic" if _use_hydroelastic else "sdf-only", + counts["finger"], + counts["nut"], + counts["bolt"], + counts["panda"], + counts["table"], + counts["skip_no_mesh"], + counts["skip_already_built"], + )