diff --git a/.gitignore b/.gitignore index 4bd3675e9df8..bbe4d2cdf505 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,8 @@ # No USD files allowed in the repo **/*.usd **/*.usda +!source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/*.usda +!source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/*.usd **/*.usdc **/*.usdz diff --git a/docs/index.rst b/docs/index.rst index 8f399b0e8b55..2edd30e1cfe3 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -91,6 +91,7 @@ Table of Contents source/setup/ecosystem source/setup/installation/index source/setup/environments + source/setup/conveyor_franka source/setup/quickstart source/setup/tutorial source/setup/demos diff --git a/docs/source/_static/conveyor_franka.jpg b/docs/source/_static/conveyor_franka.jpg new file mode 100644 index 000000000000..7ee666022966 Binary files /dev/null and b/docs/source/_static/conveyor_franka.jpg differ diff --git a/docs/source/_static/css/environment-browser.js b/docs/source/_static/css/environment-browser.js index dadfd2dca455..42751abf27ee 100644 --- a/docs/source/_static/css/environment-browser.js +++ b/docs/source/_static/css/environment-browser.js @@ -63,6 +63,9 @@ ["IsaacContrib-AutoMate-Disassembly-Direct", "rl_games", "", "", "", "tasks/automate/01053_disassembly.jpg"], ["IsaacContrib-Cartpole-Camera-Showcase-Direct", "skrl", "", "", "box_box,box_discrete,box_multidiscrete,dict_box,dict_discrete,dict_multidiscrete,tuple_box,tuple_discrete,tuple_multidiscrete"], ["IsaacContrib-Cartpole-Showcase-Direct", "skrl", "isaacsim_physx,newton_kamino,newton_mjwarp,ovphysx", "", "box_box,box_discrete,box_multidiscrete,dict_box,dict_discrete,dict_multidiscrete,discrete_box,discrete_discrete,discrete_multidiscrete,multidiscrete_box,multidiscrete_discrete,multidiscrete_multidiscrete,tuple_box,tuple_discrete,tuple_multidiscrete"], + ["IsaacContrib-Conveyor-Franka-Newton-Play-v0", "rsl_rl", "", "", "", "conveyor_franka.jpg"], + ["IsaacContrib-Conveyor-Franka-Newton-v0", "rsl_rl", "", "", "", "conveyor_franka.jpg"], + ["IsaacContrib-Conveyor-Franka-PhysX-CPU-v0", "rsl_rl", "", "", "", "conveyor_franka.jpg"], ["IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav", "rsl_rl", "", "", ""], ["IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav-ROS-Inference", "rsl_rl", "", "", ""], ["IsaacContrib-Deploy-GearAssembly-UR10e-2F140", "rsl_rl", "", "", ""], diff --git a/docs/source/api/lab/isaaclab.physics.rst b/docs/source/api/lab/isaaclab.physics.rst index d747af4083b3..12bfbb7af4b2 100644 --- a/docs/source/api/lab/isaaclab.physics.rst +++ b/docs/source/api/lab/isaaclab.physics.rst @@ -18,6 +18,8 @@ The following classes are part of the public :mod:`isaaclab.physics` API. PhysicsEvent PhysicsManager PhysxAutoCfg + SurfaceVelocitySpec + SurfaceVelocityView .. autoclass:: CallbackHandle :show-inheritance: @@ -34,3 +36,9 @@ The following classes are part of the public :mod:`isaaclab.physics` API. .. autoclass:: PhysxAutoCfg :show-inheritance: + +.. autoclass:: SurfaceVelocitySpec + :show-inheritance: + +.. autoclass:: SurfaceVelocityView + :show-inheritance: diff --git a/docs/source/api/lab_newton/isaaclab_newton.physics.rst b/docs/source/api/lab_newton/isaaclab_newton.physics.rst index 69c9f1c1b6e4..bc9121f7b999 100644 --- a/docs/source/api/lab_newton/isaaclab_newton.physics.rst +++ b/docs/source/api/lab_newton/isaaclab_newton.physics.rst @@ -36,6 +36,7 @@ KaminoPADMMSolverCfg MPMSolverCfg HydroelasticSDFCfg + SurfaceVelocity .. currentmodule:: isaaclab_newton.physics @@ -162,6 +163,13 @@ Physics Configuration :show-inheritance: :exclude-members: __init__ +Surface Velocity +---------------- + +.. autoclass:: SurfaceVelocity + :members: + :show-inheritance: + Solver Managers --------------- diff --git a/docs/source/api/lab_physx/isaaclab_physx.physics.rst b/docs/source/api/lab_physx/isaaclab_physx.physics.rst index 1830d53b08f8..e9d888238552 100644 --- a/docs/source/api/lab_physx/isaaclab_physx.physics.rst +++ b/docs/source/api/lab_physx/isaaclab_physx.physics.rst @@ -11,6 +11,9 @@ PhysxCfg PhysxBackendCfg + SurfaceVelocity + PhysxSurfaceVelocityTwist + .. currentmodule:: isaaclab_physx.physics Physics Manager @@ -33,6 +36,22 @@ Physics Configuration :show-inheritance: :exclude-members: __init__ +Surface Velocity +---------------- + +.. autoclass:: SurfaceVelocity + :members: + :show-inheritance: + +.. autoclass:: PhysxSurfaceVelocityTwist + :members: + +.. autofunction:: apply_surface_velocity_api + +.. autofunction:: compute_surface_velocity_twist + +.. autofunction:: resolve_surface_velocity_paths + Additional Public Classes ------------------------- diff --git a/docs/source/setup/conveyor_franka.rst b/docs/source/setup/conveyor_franka.rst new file mode 100644 index 000000000000..5c4188b6b628 --- /dev/null +++ b/docs/source/setup/conveyor_franka.rst @@ -0,0 +1,107 @@ +Conveyor Franka (Contrib) +========================= + +Choose between the original four-cube racetrack task and the warehouse sorting task. +Both use the same pretrained Franka policy and shared manipulation code. The sorter extends +the base environment with textured cartons, elevated returns, gravity infeeds, and color dispatch. + +.. image:: ../_static/conveyor_franka.jpg + :alt: Franka sorting colored cartons between two conveyors in a warehouse + :width: 100% + +.. list-table:: Two tasks, one policy + :header-rows: 1 + :widths: 22 24 54 + + * - Task + - Inventory + - Behavior + * - Racetrack transfer + - Four numbered cubes + - Original two closed racetracks; continuous alternating transfers. + * - Warehouse sorting + - 24 colored parcels + - Extended circulating conveyors; blue/green on one loop, orange/purple on the other. + +Both run on Newton GPU: select ``IsaacContrib-Conveyor-Franka-Newton-v0`` for the original +racetracks or ``IsaacContrib-Conveyor-Franka-Newton-Play-v0`` for sorting. The original task +also has a native PhysX CPU backend, described below. Sorting reuses the base configuration, +agent configuration, and manipulation terms; only the warehouse adds parcel-slot reassignment +and color-based dispatch. + +Run the pretrained policy +------------------------- + +Use the standard Isaac Lab installation with the Isaac Sim extra for Kit/RTX visuals: + +.. code-block:: bash + + uv run --extra isaacsim isaaclab play --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Franka-Newton-Play-v0 \ + --checkpoint https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/6.1/Isaac/IsaacLab/PretrainedCheckpoints/rsl_rl/IsaacContrib-Conveyor-Franka-Newton-v0_newtonmjwarp_none_rsl_rl.pt --num_envs 1 --device cuda:0 --viz kit --real-time \ + --kit_args=--/UJITSO/geometry=false + +The first launch downloads the referenced Omniverse assets. The warehouse uses USD-authored +materials and lighting; Kit/RTX is the intended viewer. The launch override disables experimental +geometry streaming, including saved Kit preferences, which can hide meshes updated through Fabric. +To record this view, append ``--video --video_length 1440 env.sim.physics.use_cuda_graph=False``. +Disable CUDA graphs for this Kit recording path; compact training retains its graph-enabled default. + +For compact, lightweight playback: + +.. code-block:: bash + + uv run isaaclab play --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Franka-Newton-v0 \ + --checkpoint pretrained --num_envs 8 --device cuda:0 --viz newton_gl --real-time + +The published policy is iteration 7998 of the +`conveyor training run `__. +The warehouse command uses that same checkpoint URL explicitly. Both retain 123 observations, +eight actions, 120 Hz physics, and a 60 Hz policy rate. Checkpoints remain outside the repository. + +Train or compare backends +------------------------- + +.. code-block:: bash + + uv run isaaclab train --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Franka-Newton-v0 --num_envs 256 --device cuda:0 + + uv run --extra isaacsim isaaclab play --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Franka-PhysX-CPU-v0 \ + --checkpoint /path/to/model.pt --num_envs 1 --device cpu --viz kit --real-time \ + agent.device=cpu + +Native PhysX surface velocity requires CPU simulation for this task. GPU dynamics can drop +belt contacts in the supported Isaac Sim runtime. The PhysX variant rejects CUDA devices; +its separately named pretrained artifact is not published. Policy shape compatibility does +not imply identical behavior between backends. + +Warehouse sorting +----------------- + +The original manipulation straights and adjoining 90-degree bends remain fixed. Twenty-four +40 mm cartons circulate through the extended layout. Each reset shuffles a balanced batch of +six blue, six green, six orange, and six purple cartons. Blue/green belong on the positive-Y +loop; orange/purple belong on the negative-Y loop. Four policy slots are reassigned to arriving +parcels while an active grasp retains its identity. Destination classes are supervisory metadata; +the state-based checkpoint does not recognize colors from images. + +This is a presentation and policy-reuse demonstration, not a reliably solved 24-parcel benchmark. +The unchanged checkpoint can miss grasps and reset before finishing a batch. The original +four-cube training and CPU reference configurations retain their compact layout. + +.. raw:: html + + + +`Download the warehouse preview `__. + +The task's +:download:`README <../../../source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/README.md>` +describes the USD assets, collision ownership, slot adapter, and sorting metrics. diff --git a/source/isaaclab/changelog.d/maximiliank-custom-mesh.minor.rst b/source/isaaclab/changelog.d/maximiliank-custom-mesh.minor.rst new file mode 100644 index 000000000000..e81d993575ef --- /dev/null +++ b/source/isaaclab/changelog.d/maximiliank-custom-mesh.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added :class:`~isaaclab.sim.MeshCustomCfg` for spawning meshes from authored + vertices and triangular faces. diff --git a/source/isaaclab/changelog.d/maximiliank-surface-velocity.minor.rst b/source/isaaclab/changelog.d/maximiliank-surface-velocity.minor.rst new file mode 100644 index 000000000000..4043fd658003 --- /dev/null +++ b/source/isaaclab/changelog.d/maximiliank-surface-velocity.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added backend-neutral surface-velocity descriptions and a tensorized control contract under + ``isaaclab.physics.surface_velocity``. diff --git a/source/isaaclab/isaaclab/physics/__init__.pyi b/source/isaaclab/isaaclab/physics/__init__.pyi index 92d9bfac2ff4..b31b566729f1 100644 --- a/source/isaaclab/isaaclab/physics/__init__.pyi +++ b/source/isaaclab/isaaclab/physics/__init__.pyi @@ -9,7 +9,10 @@ __all__ = [ "PhysicsManager", "PhysicsCfg", "PhysxAutoCfg", + "SurfaceVelocitySpec", + "SurfaceVelocityView", ] from .physics_manager import CallbackHandle, PhysicsEvent, PhysicsManager from .physics_manager_cfg import PhysicsCfg, PhysxAutoCfg +from .surface_velocity import SurfaceVelocitySpec, SurfaceVelocityView diff --git a/source/isaaclab/isaaclab/physics/surface_velocity.py b/source/isaaclab/isaaclab/physics/surface_velocity.py new file mode 100644 index 000000000000..f00badb6b941 --- /dev/null +++ b/source/isaaclab/isaaclab/physics/surface_velocity.py @@ -0,0 +1,179 @@ +# 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 + +"""Backend-neutral surface-velocity descriptions and control interface.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +_ENV_REGEX_NS = "{ENV_REGEX_NS}" + + +def _validate_prim_path(value: str) -> None: + """Validate one exact USD prim path or supported replicated-path template.""" + if not isinstance(value, str) or not value: + raise ValueError("Surface prim_path must be a non-empty string.") + if value.startswith(f"{_ENV_REGEX_NS}/"): + path = value[len(_ENV_REGEX_NS) :] + elif value.startswith("/"): + path = value + else: + raise ValueError(f"Surface prim_path must be absolute or start with '{_ENV_REGEX_NS}/', got {value!r}.") + if "{" in path or "}" in path: + raise ValueError(f"Surface prim_path supports only a leading '{_ENV_REGEX_NS}' placeholder, got {value!r}.") + components = path[1:].split("/") + if not components or any(not component or component in {".", ".."} for component in components): + raise ValueError(f"Surface prim_path must identify a concrete prim without empty components, got {value!r}.") + if any(any(character.isspace() for character in component) for component in components): + raise ValueError(f"Surface prim_path components must not contain whitespace, got {value!r}.") + + +def _validate_scalar(name: str, value: Any) -> float: + """Return one finite float with consistent validation errors.""" + try: + result = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"Surface {name} must be finite, got {value!r}.") from exc + if not math.isfinite(result): + raise ValueError(f"Surface {name} must be finite, got {value!r}.") + return result + + +def _validate_vector(name: str, value: tuple[float, ...], length: int, *, nonzero: bool = False) -> tuple[float, ...]: + """Return one finite vector with a stable tuple representation.""" + try: + result = tuple(float(component) for component in value) + except (TypeError, ValueError) as exc: + raise ValueError(f"Surface {name} must contain {length} finite values, got {value!r}.") from exc + if len(result) != length or not all(math.isfinite(component) for component in result): + raise ValueError(f"Surface {name} must contain {length} finite values, got {value!r}.") + if nonzero and math.sqrt(sum(component * component for component in result)) <= 1.0e-8: + raise ValueError(f"Surface {name} must be non-zero, got {value!r}.") + return result + + +@dataclass(frozen=True, slots=True) +class SurfaceVelocitySpec: + """Persistent intent for one collision surface with prescribed tangential velocity. + + The fields follow the authored conveyor model proposed for Isaac Sim while remaining independent of + Kit, OpenUSD, and any physics backend. Directions, surface normals, and the optional pivot point are + expressed in the collision prim's local frame. Runtime state such as encoder positions and applied + forces intentionally does not belong in this description. + + ``prim_path`` may contain Isaac Lab's ``{ENV_REGEX_NS}`` placeholder. A backend resolves that template + to every replicated collision prim while preserving one deterministic belt index per environment. An + absolute path addresses one unreplicated surface; replicated environments should use the placeholder. + + Args: + prim_path: Collision prim path or replicated Isaac Lab prim-path template. + velocity: Initial signed centerline surface velocity [m/s]. + enabled: Whether the surface initially transports contacting bodies. + direction: Local travel direction, or local rotation axis for a curved belt. + curved: Whether the velocity field rotates about ``pivot_point``. + pivot_point: Local rotation center for a curved belt [m]. An OpenUSD adapter may map this point to the + world transform of a prim targeted by the proposed schema's pivot relationship; it is not copied + into a schema attribute. + radius: Optional centerline radius used by backends that cannot derive it from geometry [m]. + surface_normal: Local outward normal of the carrying surface. + contact_threshold: Minimum contact-normal alignment accepted for traction. + friction_coefficient: Coulomb limit for synthetic surface traction. + """ + + prim_path: str + velocity: float = 0.0 + enabled: bool = True + direction: tuple[float, float, float] = (1.0, 0.0, 0.0) + curved: bool = False + pivot_point: tuple[float, float, float] = (0.0, 0.0, 0.0) + radius: float | None = None + surface_normal: tuple[float, float, float] = (0.0, 0.0, 1.0) + contact_threshold: float = 0.997 + friction_coefficient: float = 0.7 + + def __post_init__(self) -> None: + """Normalize immutable vectors and validate authored values.""" + _validate_prim_path(self.prim_path) + for name in ("enabled", "curved"): + if not isinstance(getattr(self, name), bool): + raise ValueError(f"Surface {name} must be a bool, got {getattr(self, name)!r}.") + + velocity = _validate_scalar("velocity", self.velocity) + contact_threshold = _validate_scalar("contact_threshold", self.contact_threshold) + friction_coefficient = _validate_scalar("friction_coefficient", self.friction_coefficient) + if not 0.0 <= contact_threshold <= 1.0: + raise ValueError(f"Surface contact_threshold must be in [0, 1], got {self.contact_threshold!r}.") + if friction_coefficient < 0.0: + raise ValueError( + f"Surface friction_coefficient must be finite and non-negative, got {self.friction_coefficient!r}." + ) + radius = None if self.radius is None else _validate_scalar("radius", self.radius) + if radius is not None and radius <= 0.0: + raise ValueError(f"Surface radius must be finite and positive when provided, got {self.radius!r}.") + + object.__setattr__(self, "velocity", velocity) + object.__setattr__(self, "direction", _validate_vector("direction", self.direction, 3, nonzero=True)) + object.__setattr__(self, "pivot_point", _validate_vector("pivot_point", self.pivot_point, 3)) + object.__setattr__( + self, "surface_normal", _validate_vector("surface_normal", self.surface_normal, 3, nonzero=True) + ) + object.__setattr__(self, "radius", radius) + object.__setattr__(self, "contact_threshold", contact_threshold) + object.__setattr__(self, "friction_coefficient", friction_coefficient) + + +@runtime_checkable +class SurfaceVelocityView(Protocol): + """Common tensorized control contract implemented by physics backends.""" + + @property + def prim_paths(self) -> tuple[str, ...]: + """Resolved collision prim paths in stable surface-index order.""" + ... + + @property + def num_surfaces(self) -> int: + """Number of resolved moving surfaces.""" + ... + + @property + def count(self) -> int: + """Alias for :attr:`num_surfaces`, matching tensor-view naming.""" + ... + + def set_velocities(self, velocities: Any, indices: Any = None) -> None: + """Set signed surface velocities [m/s] for selected surfaces.""" + ... + + def get_velocities(self, indices: Any = None, clone: bool = True) -> Any: + """Return effective surface velocities [m/s] for selected surfaces.""" + ... + + def get_commanded_velocities(self, indices: Any = None, clone: bool = True) -> Any: + """Return commanded surface velocities [m/s] before the enabled mask.""" + ... + + def set_enabled(self, flags: Any, indices: Any = None) -> None: + """Enable or disable selected surfaces without discarding their commands.""" + ... + + def get_enabled(self, indices: Any = None, clone: bool = True) -> Any: + """Return integer enabled flags for selected surfaces.""" + ... + + def get_encoder_positions(self, indices: Any = None, clone: bool = True) -> Any: + """Return integrated surface travel [m] for selected surfaces.""" + ... + + def reset(self, env_ids: Any = None) -> None: + """Clear runtime state for selected replicated environments.""" + ... + + def close(self) -> None: + """Release backend resources and callbacks.""" + ... diff --git a/source/isaaclab/isaaclab/sim/__init__.pyi b/source/isaaclab/isaaclab/sim/__init__.pyi index 0a22d87898f6..fd970f25b60b 100644 --- a/source/isaaclab/isaaclab/sim/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/__init__.pyi @@ -128,6 +128,7 @@ __all__ = [ "VisualMaterialCfg", "spawn_mesh_capsule", "spawn_mesh_cone", + "spawn_mesh_custom", "spawn_mesh_cuboid", "spawn_mesh_cylinder", "spawn_mesh_rectangle", @@ -135,6 +136,7 @@ __all__ = [ "MeshCapsuleCfg", "MeshCfg", "MeshConeCfg", + "MeshCustomCfg", "MeshCuboidCfg", "MeshCylinderCfg", "MeshRectangleCfg", @@ -343,6 +345,7 @@ from .spawners import ( MeshCapsuleCfg, MeshCfg, MeshConeCfg, + MeshCustomCfg, MeshCuboidCfg, MeshCylinderCfg, MeshRectangleCfg, @@ -387,6 +390,7 @@ from .spawners import ( spawn_light, spawn_mesh_capsule, spawn_mesh_cone, + spawn_mesh_custom, spawn_mesh_cuboid, spawn_mesh_cylinder, spawn_mesh_rectangle, diff --git a/source/isaaclab/isaaclab/sim/spawners/__init__.pyi b/source/isaaclab/isaaclab/sim/spawners/__init__.pyi index 339d1ef667ba..be3d9e26ced4 100644 --- a/source/isaaclab/isaaclab/sim/spawners/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/spawners/__init__.pyi @@ -44,6 +44,7 @@ __all__ = [ "VisualMaterialCfg", "spawn_mesh_capsule", "spawn_mesh_cone", + "spawn_mesh_custom", "spawn_mesh_cuboid", "spawn_mesh_cylinder", "spawn_mesh_rectangle", @@ -51,6 +52,7 @@ __all__ = [ "MeshCapsuleCfg", "MeshCfg", "MeshConeCfg", + "MeshCustomCfg", "MeshCuboidCfg", "MeshCylinderCfg", "MeshRectangleCfg", @@ -127,6 +129,7 @@ from .materials import ( from .meshes import ( spawn_mesh_capsule, spawn_mesh_cone, + spawn_mesh_custom, spawn_mesh_cuboid, spawn_mesh_cylinder, spawn_mesh_rectangle, @@ -134,6 +137,7 @@ from .meshes import ( MeshCapsuleCfg, MeshCfg, MeshConeCfg, + MeshCustomCfg, MeshCuboidCfg, MeshCylinderCfg, MeshRectangleCfg, diff --git a/source/isaaclab/isaaclab/sim/spawners/meshes/__init__.pyi b/source/isaaclab/isaaclab/sim/spawners/meshes/__init__.pyi index c853dca4d1d0..fd1e4ac5fdf5 100644 --- a/source/isaaclab/isaaclab/sim/spawners/meshes/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/spawners/meshes/__init__.pyi @@ -6,6 +6,7 @@ __all__ = [ "spawn_mesh_capsule", "spawn_mesh_cone", + "spawn_mesh_custom", "spawn_mesh_cuboid", "spawn_mesh_cylinder", "spawn_mesh_rectangle", @@ -13,6 +14,7 @@ __all__ = [ "MeshCapsuleCfg", "MeshCfg", "MeshConeCfg", + "MeshCustomCfg", "MeshCuboidCfg", "MeshCylinderCfg", "MeshRectangleCfg", @@ -22,6 +24,7 @@ __all__ = [ from .meshes import ( spawn_mesh_capsule, spawn_mesh_cone, + spawn_mesh_custom, spawn_mesh_cuboid, spawn_mesh_cylinder, spawn_mesh_rectangle, @@ -31,6 +34,7 @@ from .meshes_cfg import ( MeshCapsuleCfg, MeshCfg, MeshConeCfg, + MeshCustomCfg, MeshCuboidCfg, MeshCylinderCfg, MeshRectangleCfg, diff --git a/source/isaaclab/isaaclab/sim/spawners/meshes/meshes.py b/source/isaaclab/isaaclab/sim/spawners/meshes/meshes.py index f6bd7c68b4ff..7edf68ae4779 100644 --- a/source/isaaclab/isaaclab/sim/spawners/meshes/meshes.py +++ b/source/isaaclab/isaaclab/sim/spawners/meshes/meshes.py @@ -28,6 +28,56 @@ from . import meshes_cfg +@clone +def spawn_mesh_custom( + prim_path: str, + cfg: meshes_cfg.MeshCustomCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Create a USD mesh from explicitly authored vertices and triangular faces. + + This spawner is intended for custom static geometry such as terrain patches, + tracks, and fixtures that cannot be represented by the generated mesh + primitives. Face winding is preserved, so counter-clockwise faces define the + colliding side when exact triangle-mesh collision is selected. + + Args: + prim_path: Prim path or pattern at which to spawn the asset. + cfg: Custom mesh configuration. + translation: Translation relative to the parent prim [m]. + orientation: Quaternion orientation in ``(x, y, z, w)`` order. + **kwargs: Additional cloning keyword arguments. + + Returns: + The created root prim. + + Raises: + ValueError: If the vertex or face arrays are malformed, a face index is + out of range, or the collision approximation is unknown. + """ + del kwargs + vertices = np.asarray(cfg.vertices, dtype=np.float32) + faces = np.asarray(cfg.faces, dtype=np.int64) + if vertices.ndim != 2 or vertices.shape[1:] != (3,) or len(vertices) < 3: + raise ValueError(f"Custom mesh vertices must have shape (N, 3) with N >= 3, got {vertices.shape}.") + if faces.ndim != 2 or faces.shape[1:] != (3,) or len(faces) < 1: + raise ValueError(f"Custom mesh faces must have shape (M, 3) with M >= 1, got {faces.shape}.") + if np.any(faces < 0) or np.any(faces >= len(vertices)): + raise ValueError(f"Custom mesh face indices must be in [0, {len(vertices) - 1}].") + if cfg.collision_approximation not in schemas.MESH_APPROXIMATION_TOKENS: + raise ValueError( + f"Unknown mesh collision approximation {cfg.collision_approximation!r}. " + f"Valid options are: {list(schemas.MESH_APPROXIMATION_TOKENS)}" + ) + + mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False) + stage = get_current_stage() + _spawn_mesh_geom_from_mesh(prim_path, cfg, mesh, translation, orientation, stage=stage) + return stage.GetPrimAtPath(prim_path) + + @clone def spawn_mesh_sphere( prim_path: str, @@ -391,7 +441,7 @@ def _spawn_mesh_geom_from_mesh( "points": mesh.vertices, "faceVertexIndices": mesh.faces.flatten(), "faceVertexCounts": np.asarray([3] * len(mesh.faces)), - "subdivisionScheme": "bilinear", + "subdivisionScheme": getattr(cfg, "subdivision_scheme", "bilinear"), }, stage=stage, ) @@ -420,13 +470,16 @@ def _spawn_mesh_geom_from_mesh( ) elif cfg.collision_props is not None: # decide on type of collision approximation based on the mesh - if cfg.__class__.__name__ == "MeshSphereCfg": - collision_approximation = "boundingSphere" - elif cfg.__class__.__name__ == "MeshCuboidCfg": - collision_approximation = "boundingCube" - else: - # for: MeshCylinderCfg, MeshCapsuleCfg, MeshConeCfg - collision_approximation = "convexHull" + collision_approximation = getattr(cfg, "collision_approximation", None) + if collision_approximation is None: + if cfg.__class__.__name__ == "MeshSphereCfg": + collision_approximation = "boundingSphere" + elif cfg.__class__.__name__ == "MeshCuboidCfg": + collision_approximation = "boundingCube" + else: + # for: MeshCylinderCfg, MeshCapsuleCfg, MeshConeCfg + collision_approximation = "convexHull" + # apply collision approximation to mesh # note: for primitives, we use the convex hull approximation -- this should be sufficient for most cases. UsdPhysics.MeshCollisionAPI.Apply(mesh_prim).GetApproximationAttr().Set(collision_approximation) # collision properties anchor at the geometry prim diff --git a/source/isaaclab/isaaclab/sim/spawners/meshes/meshes_cfg.py b/source/isaaclab/isaaclab/sim/spawners/meshes/meshes_cfg.py index b969c8afd71f..0e107404e8b4 100644 --- a/source/isaaclab/isaaclab/sim/spawners/meshes/meshes_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/meshes/meshes_cfg.py @@ -86,6 +86,28 @@ class MeshCfg(RigidObjectSpawnerCfg, DeformableObjectSpawnerCfg): """ +@configclass +class MeshCustomCfg(MeshCfg): + """Configuration parameters for a mesh authored from vertices and triangular faces. + + See :meth:`spawn_mesh_custom` for more information. + """ + + func: Callable | str = "{DIR}.meshes:spawn_mesh_custom" + + vertices: tuple[tuple[float, float, float], ...] = MISSING + """Vertex positions [m].""" + + faces: tuple[tuple[int, int, int], ...] = MISSING + """Triangle vertex indices with counter-clockwise front-face winding.""" + + collision_approximation: str = "none" + """Mesh collision approximation name. Defaults to ``"none"`` for exact triangle-mesh collision.""" + + subdivision_scheme: Literal["none", "catmullClark", "loop", "bilinear"] = "none" + """USD subdivision scheme. Defaults to ``"none"`` to preserve the authored surface.""" + + @configclass class MeshSphereCfg(MeshCfg): """Configuration parameters for a sphere mesh prim with deformable properties. diff --git a/source/isaaclab/test/sim/test_spawn_meshes.py b/source/isaaclab/test/sim/test_spawn_meshes.py index aa017e7bb126..716bdcb75799 100644 --- a/source/isaaclab/test/sim/test_spawn_meshes.py +++ b/source/isaaclab/test/sim/test_spawn_meshes.py @@ -149,6 +149,25 @@ def test_invalid_edge_refinement(sim): cfg.func("/World/Invalid", cfg) +def test_spawn_custom_mesh(sim): + """Test spawning an exact collision mesh from authored triangle data.""" + cfg = sim_utils.MeshCustomCfg( + vertices=((-0.5, -0.5, 0.0), (0.5, -0.5, 0.0), (0.5, 0.5, 0.0), (-0.5, 0.5, 0.0)), + faces=((0, 1, 2), (0, 2, 3)), + collision_props=sim_utils.CollisionBaseCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.2, 0.4, 0.8)), + ) + prim = cfg.func("/World/CustomMesh", cfg) + + assert prim.IsValid() + mesh_prim = sim.stage.GetPrimAtPath("/World/CustomMesh/geometry/mesh") + assert mesh_prim.GetPrimTypeInfo().GetTypeName() == "Mesh" + assert mesh_prim.GetAttribute("faceVertexIndices").Get() == [0, 1, 2, 0, 2, 3] + assert mesh_prim.GetAttribute("subdivisionScheme").Get() == "none" + assert mesh_prim.GetAttribute("physics:approximation").Get() == "none" + assert not mesh_prim.GetAttribute("primvars:displayColor").HasAuthoredValue() + + """ Physics properties. """ diff --git a/source/isaaclab/test/sim/test_surface_velocity.py b/source/isaaclab/test/sim/test_surface_velocity.py new file mode 100644 index 000000000000..a043d3a6cfab --- /dev/null +++ b/source/isaaclab/test/sim/test_surface_velocity.py @@ -0,0 +1,66 @@ +# 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 + +"""Focused tests for backend-neutral surface-velocity descriptions.""" + +from __future__ import annotations + +import pytest + +from isaaclab.physics import SurfaceVelocitySpec + + +def test_surface_velocity_spec_preserves_authored_semantics() -> None: + """The shared description carries schema-aligned fields without backend imports.""" + spec = SurfaceVelocitySpec( + prim_path="{ENV_REGEX_NS}/Belt/Curve", + velocity=-0.35, + enabled=False, + direction=(0, 0, -1), + curved=True, + pivot_point=(0.58, 0.51, 0.0), + radius=0.24, + surface_normal=(0, 0, 1), + contact_threshold=0.997, + friction_coefficient=0.5, + ) + + assert spec.prim_path == "{ENV_REGEX_NS}/Belt/Curve" + assert spec.velocity == -0.35 + assert spec.enabled is False + assert spec.direction == (0.0, 0.0, -1.0) + assert spec.curved is True + assert spec.pivot_point == (0.58, 0.51, 0.0) + assert spec.radius == 0.24 + assert spec.surface_normal == (0.0, 0.0, 1.0) + assert spec.contact_threshold == 0.997 + assert spec.friction_coefficient == 0.5 + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"prim_path": "relative/Belt"}, "prim_path"), + ({"prim_path": "/"}, "prim_path"), + ({"prim_path": "/World//Belt"}, "prim_path"), + ({"prim_path": "/World/Belt/"}, "prim_path"), + ({"prim_path": "{ENV_REGEX_NS}/"}, "prim_path"), + ({"prim_path": "{UNKNOWN_NS}/Belt"}, "prim_path"), + ({"prim_path": "/World/{ENV_REGEX_NS}/Belt"}, "prim_path"), + ({"prim_path": "/World/Belt", "velocity": None}, "velocity"), + ({"prim_path": "/World/Belt", "velocity": float("nan")}, "velocity"), + ({"prim_path": "/World/Belt", "enabled": 1}, "enabled"), + ({"prim_path": "/World/Belt", "curved": "yes"}, "curved"), + ({"prim_path": "/World/Belt", "direction": (0.0, 0.0, 0.0)}, "direction"), + ({"prim_path": "/World/Belt", "surface_normal": (0.0, 0.0)}, "surface_normal"), + ({"prim_path": "/World/Belt", "radius": 0.0}, "radius"), + ({"prim_path": "/World/Belt", "contact_threshold": 1.1}, "contact_threshold"), + ({"prim_path": "/World/Belt", "friction_coefficient": -0.1}, "friction_coefficient"), + ], +) +def test_surface_velocity_spec_rejects_invalid_authored_values(kwargs: dict, message: str) -> None: + """Invalid persistent intent fails before any physics lifecycle is registered.""" + with pytest.raises(ValueError, match=message): + SurfaceVelocitySpec(**kwargs) diff --git a/source/isaaclab_newton/changelog.d/maximiliank-conveyor-substep-callback.rst b/source/isaaclab_newton/changelog.d/maximiliank-conveyor-substep-callback.rst new file mode 100644 index 000000000000..c111b8a747c8 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/maximiliank-conveyor-substep-callback.rst @@ -0,0 +1,8 @@ +Added +^^^^^ + +* Added lifecycle-safe Newton manager callbacks for binding model-specific resources before CUDA graph capture + and applying contact-force feedback after each solver substep. +* Added reusable, batched surface-velocity contact forces under ``isaaclab_newton.physics.surface_velocity``. +* Exposed ``NewtonCollisionPipelineCfg.include_static_kinematic_pairs`` to exclude contacts between + immovable shapes without affecting contacts with dynamic bodies. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi index 39a2887807d7..21aa12bb3d15 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi @@ -30,6 +30,7 @@ __all__ = [ "NewtonShapeCfg", "NewtonSoftContactCfg", "NewtonSolverCfg", + "SurfaceVelocity", "NewtonVBDManager", "VBDSolverCfg", "NewtonXPBDManager", @@ -65,6 +66,7 @@ from .newton_manager_cfg import ( NewtonSoftContactCfg, NewtonSolverCfg, ) +from .surface_velocity import SurfaceVelocity from .vbd_manager import NewtonVBDManager from .vbd_manager_cfg import VBDSolverCfg from .xpbd_manager import NewtonXPBDManager 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 8bb8b75cae0f..a13c4a84c21e 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_collision_cfg.py @@ -106,6 +106,13 @@ class NewtonCollisionPipelineCfg: Defaults to ``"explicit"`` (same as Newton's default when ``broad_phase=None``). """ + include_static_kinematic_pairs: bool = True + """Whether to generate contacts between two immovable shapes. + + Set to ``False`` to exclude static-static, static-kinematic, and kinematic-kinematic + pairs. Defaults to ``True``, matching Newton's default. + """ + reduce_contacts: bool = True """Whether to reduce contacts for mesh-mesh collisions. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 0d01ef42dc7a..d7140b6fc77e 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -482,6 +482,13 @@ class NewtonManager(PhysicsManager): _post_actuator_callbacks: list[Callable[[], None]] = [] # In-graph hooks invoked immediately before every solver substep. _state_force_callbacks: list[Callable[[State], None]] = [] + # In-graph hooks invoked after every solver substep, before the post-step + # state is swapped into the active buffer and its external forces cleared. + _post_solver_substep_callbacks: list[Callable[[SolverBase, Contacts | None, State, float], None]] = [] + # Lifecycle hooks invoked after the solver and contact buffers are created, + # but before any CUDA graph is captured. Scene-owned force systems use this + # seam to bind model-specific buffers and register their in-graph callbacks. + _solver_init_callbacks: list[Callable[[Model, Contacts | None], None]] = [] # In-graph hooks invoked after the last solver substep and before sensors, # in registration order. Articulations with non-identity ordering register # their backend-to-user state republish kernels here so the reorders are @@ -792,6 +799,8 @@ def clear(cls): NewtonManager._adapter = None NewtonManager._post_actuator_callbacks = [] NewtonManager._state_force_callbacks = [] + NewtonManager._post_solver_substep_callbacks = [] + NewtonManager._solver_init_callbacks = [] NewtonManager._post_step_callbacks = [] # Set by an articulation that took the ``use_newton_actuators=True`` # branch in ``_process_actuators_cfg``. Together with the adapter @@ -1578,6 +1587,12 @@ def initialize_solver(cls) -> None: ) cls._initialize_contacts() + # Scene-owned systems that consume solver/contact buffers must bind + # here: both resources now exist, while CUDA graph capture has not yet + # started. The callbacks intentionally persist across hard resets and + # rebind to each re-finalized model. + cls._run_solver_init_callbacks() + # Picking callbacks must be registered after the concrete solver has # published its force-input capability, but before CUDA graph capture. sim = PhysicsManager._sim @@ -1645,6 +1660,8 @@ def _run_solver_substeps(cls, contacts) -> None: for callback in cls._state_force_callbacks: callback(backend.state_0) cls._step_solver(backend.state_0, backend.state_0, backend.control, contacts, cls._solver_dt) + for cb in cls._post_solver_substep_callbacks: + cb(cls._solver, contacts, backend.state_0, cls._solver_dt) backend.state_0.clear_forces() if collide_mid_loop and (i + 1) % collide_every == 0 and i + 1 < cls._num_substeps: cls._collision_pipeline.collide(backend.state_0, contacts) @@ -1655,6 +1672,8 @@ def _run_solver_substeps(cls, contacts) -> None: for callback in cls._state_force_callbacks: callback(backend.state_0) cls._step_solver(backend.state_0, backend.state_1, backend.control, contacts, cls._solver_dt) + for cb in cls._post_solver_substep_callbacks: + cb(cls._solver, contacts, backend.state_1, cls._solver_dt) if need_copy_on_last and i == cls._num_substeps - 1: backend.state_0.assign(backend.state_1) else: @@ -1871,6 +1890,98 @@ def register_state_force_callback(cls, callback: Callable[[State], None]) -> Non return NewtonManager._state_force_callbacks.append(callback) + @classmethod + def unregister_state_force_callback(cls, callback: Callable[[State], None]) -> None: + """Remove a previously registered state-force callback. + + Removing a callback that was never registered or was already removed is + a safe no-op. This lets scene-owned systems release bound-method + references before the global Newton manager is cleared. + + Args: + callback: Previously registered callback. + """ + with contextlib.suppress(ValueError): + NewtonManager._state_force_callbacks.remove(callback) + + @classmethod + def unregister_post_actuator_callback(cls, callback: Callable[[], None]) -> None: + """Remove a previously registered post-actuator callback. + + Removing a callback that was never registered or was already removed is + a safe no-op. This lets scene-owned systems release bound-method + references before the global Newton manager is cleared. + + Args: + callback: Previously registered callback. + """ + with contextlib.suppress(ValueError): + cls._post_actuator_callbacks.remove(callback) + + @classmethod + def register_post_solver_substep_callback( + cls, callback: Callable[[SolverBase, Contacts | None, State, float], None] + ) -> None: + """Append a hook invoked immediately after every solver substep. + + The callback receives the active solver, contacts, post-substep state, + and solver timestep [s]. It runs before double-buffered state swapping + and before external forces are cleared, matching Newton's native + per-substep force-feedback loop. Callbacks must be graph-safe. + + Args: + callback: Function called after each solver substep. + """ + if callback in NewtonManager._post_solver_substep_callbacks: + return + cls._post_solver_substep_callbacks.append(callback) + + @classmethod + def register_solver_init_callback(cls, callback: Callable[[Model, Contacts | None], None]) -> None: + """Register a callback that binds resources before CUDA graph capture. + + The callback runs after every solver/contact initialization, including + hard resets, and before any simulation graph is captured. It may + allocate model-specific buffers and register graph-safe step callbacks. + + Args: + callback: Function receiving the active model and contact buffer. + """ + if callback in NewtonManager._solver_init_callbacks: + return + NewtonManager._solver_init_callbacks.append(callback) + + @classmethod + def _run_solver_init_callbacks(cls) -> None: + """Bind scene-owned solver resources before graph capture.""" + for callback in tuple(NewtonManager._solver_init_callbacks): + callback(cls.backend.model, cls._contacts) + + @classmethod + def unregister_solver_init_callback(cls, callback: Callable[[Model, Contacts | None], None]) -> None: + """Remove a previously registered solver-initialization callback. + + Removing a callback that was never registered or was already removed is + a safe no-op. + + Args: + callback: Previously registered callback. + """ + with contextlib.suppress(ValueError): + NewtonManager._solver_init_callbacks.remove(callback) + + @classmethod + def unregister_post_solver_substep_callback( + cls, callback: Callable[[SolverBase, Contacts | None, State, float], None] + ) -> None: + """Remove a previously registered post-solver-substep callback. + + Args: + callback: Previously registered callback. + """ + with contextlib.suppress(ValueError): + cls._post_solver_substep_callbacks.remove(callback) + @classmethod def register_post_step_callback(cls, callback: Callable[[], None]) -> None: """Append a hook to the list invoked after the last solver substep on every step. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/surface_velocity.py b/source/isaaclab_newton/isaaclab_newton/physics/surface_velocity.py new file mode 100644 index 000000000000..3c8109bf9ca1 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/surface_velocity.py @@ -0,0 +1,1182 @@ +# 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 + +"""Batched contact-force surface velocities for Newton physics. + +The driver reads solver-reported normal contact forces, computes a Coulomb-limited force that +drives each transported body's contact points toward their conveyor velocity fields, and applies +the resulting body wrenches on the following physics solve. +""" + +from __future__ import annotations + +import re +from collections.abc import Sequence +from typing import Any + +import numpy as np +import warp as wp + +from isaaclab.physics import PhysicsEvent, SurfaceVelocitySpec + +from .newton_manager import NewtonManager + +_VELOCITY_FIELD_TYPE_CONSTANT = 0 +_VELOCITY_FIELD_TYPE_PIVOT = 1 +_SAME_NORMAL_THRESHOLD = 0.999 + + +@wp.struct +class Vec3Pair: + """Two orthonormal vectors spanning a contact tangent plane.""" + + v0: wp.vec3 + v1: wp.vec3 + + +@wp.func +def compute_basis_vectors(direction: wp.vec3) -> Vec3Pair: + """Return the reference conveyor's tangent basis for a unit direction.""" + basis = Vec3Pair() + if wp.abs(direction[1]) <= 0.9999: + basis.v0 = wp.normalize(wp.vec3(direction[2], 0.0, -direction[0])) + basis.v1 = wp.vec3( + direction[1] * basis.v0[2], + (direction[2] * basis.v0[0]) - (direction[0] * basis.v0[2]), + -direction[1] * basis.v0[0], + ) + else: + basis.v0 = wp.vec3(1.0, 0.0, 0.0) + basis.v1 = wp.normalize(wp.vec3(0.0, direction[2], -direction[1])) + return basis + + +@wp.func +def compute_point_impulse( + normal: wp.vec3, + normal_impulse: wp.float32, + current_vel: wp.vec3, + target_vel: wp.vec3, + response_linear: wp.float32, + inv_inertia_world: wp.mat33, + center_of_mass_to_point: wp.vec3, + friction_coefficient: wp.float32, + mass_splitting_scale: wp.float32, +) -> wp.vec3: + """Compute a Coulomb-clamped tangential impulse using point effective mass.""" + rel_vel = target_vel - current_vel + basis = compute_basis_vectors(normal) + + r_cross_t0 = wp.cross(center_of_mass_to_point, basis.v0) + r_cross_t1 = wp.cross(center_of_mass_to_point, basis.v1) + k00 = response_linear + wp.dot(r_cross_t0, wp.mul(inv_inertia_world, r_cross_t0)) + k11 = response_linear + wp.dot(r_cross_t1, wp.mul(inv_inertia_world, r_cross_t1)) + k01 = wp.dot(r_cross_t0, wp.mul(inv_inertia_world, r_cross_t1)) + det = (k00 * k11) - (k01 * k01) + + i0 = wp.float32(0.0) + i1 = wp.float32(0.0) + if det > 0.0: + v0 = wp.dot(basis.v0, rel_vel) + v1 = wp.dot(basis.v1, rel_vel) + i0 = ((k11 * v0) - (k01 * v1)) * mass_splitting_scale / det + i1 = ((k00 * v1) - (k01 * v0)) * mass_splitting_scale / det + + friction_impulse_max = normal_impulse * friction_coefficient + zero_err_magn = wp.sqrt((i0 * i0) + (i1 * i1)) + impulse_magn = wp.min(friction_impulse_max, zero_err_magn) + if zero_err_magn > 0.0: + ratio = impulse_magn / zero_err_magn + else: + ratio = 0.0 + return (basis.v0 * (i0 * ratio)) + (basis.v1 * (i1 * ratio)) + + +@wp.func +def compute_point_force( + dt: wp.float32, + inverse_dt: wp.float32, + com_world: wp.vec3, + body_inverse_mass: wp.float32, + body_inverse_inertia_world: wp.mat33, + body_linear_velocity: wp.vec3, + body_angular_velocity: wp.vec3, + contact_position: wp.vec3, + contact_normal: wp.vec3, + contact_force: wp.float32, + mass_splitting_scale: wp.float32, + target_vel: wp.vec3, + friction_coefficient: wp.float32, +) -> wp.spatial_vector: + """Compute force and torque at one conveyor contact.""" + contact_impulse = contact_force * dt + center_of_mass_to_point = contact_position - com_world + current_point_vel = body_linear_velocity + wp.cross(body_angular_velocity, center_of_mass_to_point) + + tangential_impulse = compute_point_impulse( + contact_normal, + contact_impulse, + current_point_vel, + target_vel, + body_inverse_mass, + body_inverse_inertia_world, + center_of_mass_to_point, + friction_coefficient, + mass_splitting_scale, + ) + + force = tangential_impulse * inverse_dt + torque = wp.cross(center_of_mass_to_point, force) + return wp.spatial_vector(force, torque) + + +@wp.struct +class BeltContact: + """Reduced contact data consumed by the conveyor force kernels.""" + + valid: wp.int32 + body: wp.int32 + conveyor: wp.int32 + point: wp.vec3 + normal: wp.vec3 + normal_force: wp.float32 + next_body_contact: wp.int32 + + +@wp.kernel +def _extract_linear_force(spatial_force: wp.array[wp.spatial_vector], force: wp.array[wp.vec3]): + contact_id = wp.tid() + force[contact_id] = wp.spatial_top(spatial_force[contact_id]) + + +@wp.kernel +def _classify_contacts( + contact_count: wp.array[wp.int32], + shape0: wp.array[wp.int32], + shape1: wp.array[wp.int32], + normal: wp.array[wp.vec3], + point0: wp.array[wp.vec3], + point1: wp.array[wp.vec3], + contact_force: wp.array[wp.vec3], + shape_body: wp.array[wp.int32], + shape_conveyor: wp.array[wp.int32], + body_is_tracked: wp.array[wp.int32], + body_q: wp.array[wp.transform], + conveyor_surface_normal: wp.array[wp.vec3], + conveyor_threshold: wp.array[wp.float32], + contacts_out: wp.array[BeltContact], + body_contact_head: wp.array[wp.int32], +): + contact_id = wp.tid() + result = BeltContact() + result.valid = 0 + result.next_body_contact = -1 + + if contact_id < contact_count[0]: + contact_shape0 = shape0[contact_id] + contact_shape1 = shape1[contact_id] + if contact_shape0 >= 0 and contact_shape1 >= 0: + conveyor0 = shape_conveyor[contact_shape0] + conveyor1 = shape_conveyor[contact_shape1] + contact_normal = normal[contact_id] + + body = wp.int32(-1) + conveyor = wp.int32(-1) + local_point = wp.vec3() + normal_toward_body = wp.vec3() + if conveyor0 >= 0 and conveyor1 < 0: + conveyor = conveyor0 + body = shape_body[contact_shape1] + local_point = point1[contact_id] + normal_toward_body = contact_normal + elif conveyor1 >= 0 and conveyor0 < 0: + conveyor = conveyor1 + body = shape_body[contact_shape0] + local_point = point0[contact_id] + normal_toward_body = -contact_normal + + if body >= 0 and conveyor >= 0 and body_is_tracked[body] != 0: + alignment = wp.dot(normal_toward_body, conveyor_surface_normal[conveyor]) + normal_force = wp.abs(wp.dot(contact_force[contact_id], contact_normal)) + if alignment >= conveyor_threshold[conveyor] and normal_force > 0.0: + result.valid = 1 + result.body = body + result.conveyor = conveyor + result.point = wp.transform_point(body_q[body], local_point) + result.normal = normal_toward_body + result.normal_force = normal_force + result.next_body_contact = wp.atomic_exch(body_contact_head, body, contact_id) + + contacts_out[contact_id] = result + + +@wp.kernel +def _prepare_contact_patches( + contacts: wp.array[BeltContact], + body_contact_head: wp.array[wp.int32], + body_q: wp.array[wp.transform], + body_com: wp.array[wp.vec3], + contact_patch_head: wp.array[wp.int32], + adjusted_normal_force: wp.array[wp.float32], + mass_splitting_scale: wp.array[wp.float32], +): + """Correlate contacts by normal and normalize loads across overlapping sections.""" + body_id = wp.tid() + reference_point = wp.transform_point(body_q[body_id], body_com[body_id]) + patch_contact_id = body_contact_head[body_id] + + while patch_contact_id >= 0: + if contact_patch_head[patch_contact_id] < 0: + patch_contact = contacts[patch_contact_id] + basis = compute_basis_vectors(patch_contact.normal) + point_count = wp.int32(0) + first_conveyor = patch_contact.conveyor + spans_multiple_conveyors = wp.int32(0) + patch_force_sum = wp.float32(0.0) + min0 = wp.float32(1.0e30) + max0 = wp.float32(-1.0e30) + min1 = wp.float32(1.0e30) + max1 = wp.float32(-1.0e30) + + contact_id = body_contact_head[body_id] + while contact_id >= 0: + contact = contacts[contact_id] + if ( + contact_patch_head[contact_id] < 0 + and wp.dot(patch_contact.normal, contact.normal) > _SAME_NORMAL_THRESHOLD + ): + contact_patch_head[contact_id] = patch_contact_id + point_count += 1 + if contact.conveyor != first_conveyor: + spans_multiple_conveyors = 1 + patch_force_sum += contact.normal_force + + delta = contact.point - reference_point + projection0 = wp.dot(basis.v0, delta) + projection1 = wp.dot(basis.v1, delta) + min0 = wp.min(min0, projection0) + max0 = wp.max(max0, projection0) + min1 = wp.min(min1, projection1) + max1 = wp.max(max1, projection1) + contact_id = contact.next_body_contact + + splitting_scale = 1.0 / wp.float32(point_count) + if point_count == 1 or spans_multiple_conveyors == 0: + contact_id = body_contact_head[body_id] + while contact_id >= 0: + contact = contacts[contact_id] + if contact_patch_head[contact_id] == patch_contact_id: + adjusted_normal_force[contact_id] = contact.normal_force + mass_splitting_scale[contact_id] = splitting_scale + contact_id = contact.next_body_contact + else: + kernel_radius = 0.25 * ((max0 - min0) + (max1 - min1)) + kernel_radius_sqr = kernel_radius * kernel_radius + if kernel_radius > 0.0: + point_force_weight_sum = wp.float32(0.0) + contact_id = body_contact_head[body_id] + while contact_id >= 0: + contact = contacts[contact_id] + if contact_patch_head[contact_id] == patch_contact_id: + density = wp.float32(1.0) + other_contact_id = body_contact_head[body_id] + while other_contact_id >= 0: + other_contact = contacts[other_contact_id] + if ( + other_contact_id != contact_id + and contact_patch_head[other_contact_id] == patch_contact_id + ): + delta = contact.point - other_contact.point + projected_delta = delta - ( + wp.dot(delta, patch_contact.normal) * patch_contact.normal + ) + density += wp.exp(-0.5 * wp.length_sq(projected_delta) / kernel_radius_sqr) + other_contact_id = other_contact.next_body_contact + + weight = 1.0 / density + adjusted_normal_force[contact_id] = weight + mass_splitting_scale[contact_id] = splitting_scale + point_force_weight_sum += weight + contact_id = contact.next_body_contact + + force_per_weight = patch_force_sum / point_force_weight_sum + contact_id = body_contact_head[body_id] + while contact_id >= 0: + contact = contacts[contact_id] + if contact_patch_head[contact_id] == patch_contact_id: + adjusted_normal_force[contact_id] *= force_per_weight + contact_id = contact.next_body_contact + else: + adjusted_force = patch_force_sum / wp.float32(point_count) + contact_id = body_contact_head[body_id] + while contact_id >= 0: + contact = contacts[contact_id] + if contact_patch_head[contact_id] == patch_contact_id: + adjusted_normal_force[contact_id] = adjusted_force + mass_splitting_scale[contact_id] = splitting_scale + contact_id = contact.next_body_contact + + patch_contact_id = contacts[patch_contact_id].next_body_contact + + +@wp.kernel +def _accumulate_forces( + dt: wp.float32, + contacts: wp.array[BeltContact], + body_q: wp.array[wp.transform], + body_qd: wp.array[wp.spatial_vector], + body_com: wp.array[wp.vec3], + body_inv_mass: wp.array[wp.float32], + body_inv_inertia: wp.array[wp.mat33], + adjusted_normal_force: wp.array[wp.float32], + mass_splitting_scale: wp.array[wp.float32], + conveyor_field_type: wp.array[wp.int32], + conveyor_direction: wp.array[wp.vec3], + conveyor_pivot_point: wp.array[wp.vec3], + conveyor_radius: wp.array[wp.float32], + conveyor_effective_velocity: wp.array[wp.float32], + conveyor_friction: wp.array[wp.float32], + velocity_scale: wp.array[wp.float32], + body_force: wp.array[wp.spatial_vector], +): + contact_id = wp.tid() + contact = contacts[contact_id] + if contact.valid == 0: + return + + splitting_scale = mass_splitting_scale[contact_id] + if splitting_scale <= 0.0: + return + + conveyor = contact.conveyor + effective_velocity = conveyor_effective_velocity[conveyor] * velocity_scale[0] + if conveyor_field_type[conveyor] == _VELOCITY_FIELD_TYPE_CONSTANT: + target_velocity = conveyor_direction[conveyor] * effective_velocity + else: + angular_velocity = conveyor_direction[conveyor] * (effective_velocity / conveyor_radius[conveyor]) + target_velocity = wp.cross(angular_velocity, contact.point - conveyor_pivot_point[conveyor]) + + pose = body_q[contact.body] + center_of_mass = wp.transform_point(pose, body_com[contact.body]) + rotation = wp.quat_to_matrix(wp.transform_get_rotation(pose)) + inverse_inertia_world = rotation * body_inv_inertia[contact.body] * wp.transpose(rotation) + velocity = body_qd[contact.body] + + force = compute_point_force( + dt, + 1.0 / dt, + center_of_mass, + body_inv_mass[contact.body], + inverse_inertia_world, + wp.spatial_top(velocity), + wp.spatial_bottom(velocity), + contact.point, + contact.normal, + adjusted_normal_force[contact_id], + splitting_scale, + target_velocity, + conveyor_friction[conveyor], + ) + wp.atomic_add(body_force, contact.body, force) + + +@wp.kernel +def _advance_startup_scale( + dt: wp.float32, duration: wp.float32, elapsed: wp.array[wp.float32], scale: wp.array[wp.float32] +): + elapsed[0] += dt + scale[0] = wp.min(1.0, elapsed[0] / duration) + + +@wp.kernel +def _integrate_encoders( + dt: wp.float32, + effective_velocity: wp.array[wp.float32], + position: wp.array[wp.float32], +): + conveyor_id = wp.tid() + position[conveyor_id] += dt * effective_velocity[conveyor_id] + + +@wp.kernel +def _update_effective_velocities( + commanded_velocity: wp.array[wp.float32], + enabled: wp.array[wp.int32], + effective_velocity: wp.array[wp.float32], +): + conveyor_id = wp.tid() + effective_velocity[conveyor_id] = commanded_velocity[conveyor_id] * wp.float32(enabled[conveyor_id]) + + +@wp.kernel +def _gather_float_values( + source: wp.array[wp.float32], + indices: wp.array[wp.int32], + values: wp.array[wp.float32], +): + output_id = wp.tid() + values[output_id] = source[indices[output_id]] + + +@wp.kernel +def _gather_int_values( + source: wp.array[wp.int32], + indices: wp.array[wp.int32], + values: wp.array[wp.int32], +): + output_id = wp.tid() + values[output_id] = source[indices[output_id]] + + +@wp.kernel +def _add_body_force(dst: wp.array[wp.spatial_vector], src: wp.array[wp.spatial_vector]): + body_id = wp.tid() + dst[body_id] = dst[body_id] + src[body_id] + + +@wp.kernel +def _clear_selected_body_forces( + body_world: wp.array[wp.int32], + world_mask: wp.array[wp.bool], + body_force: wp.array[wp.spatial_vector], +): + body_id = wp.tid() + world_id = body_world[body_id] + if world_id >= 0 and world_mask[world_id]: + body_force[body_id] = wp.spatial_vector() + + +@wp.kernel +def _clear_selected_encoders( + conveyor_world: wp.array[wp.int32], + world_mask: wp.array[wp.bool], + encoder_position: wp.array[wp.float32], +): + conveyor_id = wp.tid() + world_id = conveyor_world[conveyor_id] + if world_mask[world_id]: + encoder_position[conveyor_id] = 0.0 + + +def _require_buffer_length(name: str, buffer: Any, expected: int) -> None: + """Reject missing or mis-sized buffers before a Warp launch can access them.""" + actual = len(buffer) if buffer is not None else 0 + if actual != expected: + raise RuntimeError(f"Conveyor force buffer {name!r} has length {actual}, expected {expected}.") + + +def _as_numpy(values: Any) -> np.ndarray: + """Convert supported tensor-like values to a host NumPy array.""" + if hasattr(values, "detach"): + return values.detach().cpu().numpy() + if hasattr(values, "numpy") and not isinstance(values, np.ndarray): + return values.numpy() + return np.asarray(values) + + +def _world_vector(transform_values: np.ndarray, local_vector: tuple[float, float, float]) -> wp.vec3: + """Rotate a local vector into world space and normalize it.""" + transform = wp.transform( + wp.vec3(*(float(value) for value in transform_values[:3])), + wp.quat(*(float(value) for value in transform_values[3:])), + ) + rotated = wp.transform_vector(transform, wp.vec3(*local_vector)) + values = np.asarray([float(rotated[index]) for index in range(3)], dtype=np.float32) + norm = float(np.linalg.norm(values)) + if norm <= 1.0e-8: + raise ValueError(f"Conveyor direction or surface normal must be non-zero, got {local_vector}.") + values /= norm + return wp.vec3(*(float(value) for value in values)) + + +def _world_point(transform_values: np.ndarray, local_point: tuple[float, float, float]) -> wp.vec3: + """Transform a local point into world space.""" + transform = wp.transform( + wp.vec3(*(float(value) for value in transform_values[:3])), + wp.quat(*(float(value) for value in transform_values[3:])), + ) + point = wp.transform_point(transform, wp.vec3(*local_point)) + return wp.vec3(*(float(point[index]) for index in range(3))) + + +def _resolve_belt_prim_path(prim_path: str, env_path_format: str, world_id: int) -> str: + """Resolve one replicated conveyor template to an exact Newton shape label.""" + if prim_path.startswith("{ENV_REGEX_NS}/"): + return prim_path.format(ENV_REGEX_NS=env_path_format.format(world_id)) + return prim_path + + +def _shape_belongs_to_prim(shape_label: str, prim_path: str) -> bool: + """Return whether a Newton collision-shape label is the prim or one of its descendants.""" + return shape_label == prim_path or shape_label.startswith(f"{prim_path.rstrip('/')}/") + + +def _validate_env_path_format(env_path_format: str) -> None: + """Validate the concrete per-world path format derived from the scene cloner.""" + if not isinstance(env_path_format, str) or env_path_format.count("{}") != 1: + raise ValueError(f"Conveyor env_path_format must contain exactly one '{{}}', got {env_path_format!r}.") + if not env_path_format.startswith("/"): + raise ValueError(f"Conveyor env_path_format must be absolute, got {env_path_format!r}.") + + +def _validate_newton_surface_specs(surface_specs: Sequence[SurfaceVelocitySpec]) -> None: + """Validate Newton-specific requirements before registering lifecycle callbacks.""" + if not surface_specs: + raise ValueError("At least one surface-velocity specification is required.") + prim_paths = [spec.prim_path for spec in surface_specs] + if len(set(prim_paths)) != len(prim_paths): + raise ValueError(f"Conveyor prim paths must be unique, got {prim_paths}.") + for index, path in enumerate(prim_paths): + for other in prim_paths[index + 1 :]: + if _shape_belongs_to_prim(path, other) or _shape_belongs_to_prim(other, path): + raise ValueError(f"Conveyor prim paths must not be ancestors of one another: {path!r}, {other!r}.") + for spec in surface_specs: + if spec.curved and spec.radius is None: + raise ValueError(f"Newton requires an explicit positive radius for curved belt {spec.prim_path!r}.") + + +class SurfaceVelocity: + """Own a contact-force surface-velocity pipeline across the Newton lifecycle. + + The driver is created after the simulation context but before its first + reset. It requests solved contact forces before model finalization, then + binds model-specific buffers after solver initialization and before CUDA + graph capture. A hard simulation reset transparently replaces that binding + with buffers for the re-finalized model. Resolved belt indices use deterministic + environment-major ordering across those rebuilds. + """ + + def __init__( + self, + num_envs: int, + surface_specs: Sequence[SurfaceVelocitySpec], + *, + body_pattern: str, + body_count_per_env: int | None = None, + startup_duration_s: float = 1.0, + env_path_format: str = "/World/envs/env_{}", + ) -> None: + """Register the force pipeline for the next Newton model initialization. + + Args: + num_envs: Number of replicated simulation environments. + surface_specs: Authored surface descriptions in stable within-environment order. + body_pattern: Regular expression selecting bodies that receive surface traction. + body_count_per_env: Expected selected body count per environment, or ``None``. + startup_duration_s: Duration of the initial traction ramp [s]. + env_path_format: Format string resolving one exact environment root from its integer world index. + """ + surface_specs = tuple(surface_specs) + if not all(isinstance(spec, SurfaceVelocitySpec) for spec in surface_specs): + raise TypeError("Every surface specification must be a SurfaceVelocitySpec.") + _validate_newton_surface_specs(surface_specs) + if not isinstance(num_envs, int) or isinstance(num_envs, bool) or num_envs <= 0: + raise ValueError(f"num_envs must be a positive integer, got {num_envs!r}.") + if num_envs > 1 and any(not spec.prim_path.startswith("{ENV_REGEX_NS}/") for spec in surface_specs): + raise ValueError( + "Replicated conveyor environments require every belt prim_path to start with '{ENV_REGEX_NS}/'." + ) + if not np.isfinite(startup_duration_s) or startup_duration_s <= 0.0: + raise ValueError(f"Conveyor startup duration must be positive, got {startup_duration_s}.") + _validate_env_path_format(env_path_format) + if not isinstance(body_pattern, str): + raise ValueError(f"body_pattern must be a regular-expression string, got {body_pattern!r}.") + try: + re.compile(body_pattern) + except re.error as exc: + raise ValueError(f"Invalid body pattern: {body_pattern!r}.") from exc + if body_count_per_env is not None and ( + not isinstance(body_count_per_env, int) or isinstance(body_count_per_env, bool) or body_count_per_env < 0 + ): + raise ValueError("body_count_per_env must be a non-negative integer or None.") + self._binding: _SurfaceVelocityBinding | None = None + self._closed = False + self._num_envs = num_envs + self._surface_specs = surface_specs + self._binding_kwargs = { + "num_envs": num_envs, + "surface_specs": self._surface_specs, + "startup_duration_s": startup_duration_s, + "env_path_format": env_path_format, + "body_pattern": body_pattern, + "body_count_per_env": body_count_per_env, + } + self._model_init_handle = NewtonManager.register_callback( + self._request_contact_forces, + PhysicsEvent.MODEL_INIT, + name="surface_velocity_contact_attribute", + ) + try: + NewtonManager.register_solver_init_callback(self._bind_solver) + except Exception: + self._model_init_handle.deregister() + raise + + def _require_binding(self) -> _SurfaceVelocityBinding: + """Return the current binding or fail before solver initialization.""" + binding = self._binding + if binding is None: + raise RuntimeError("Surface velocity is not bound to an initialized Newton solver.") + return binding + + @property + def specs(self) -> tuple[SurfaceVelocitySpec, ...]: + """Authored surface descriptions in stable within-environment order.""" + return self._surface_specs + + @property + def surfaces_per_env(self) -> int: + """Number of authored surfaces in each replicated environment.""" + return len(self._surface_specs) + + @property + def num_surfaces(self) -> int: + """Total number of resolved surfaces across all environments.""" + return self._num_envs * self.surfaces_per_env + + @property + def count(self) -> int: + """Alias for :attr:`num_surfaces`, matching tensor-view naming.""" + return self.num_surfaces + + @property + def initialized(self) -> bool: + """Whether the driver is bound to the active Newton solver.""" + return self._binding is not None + + @property + def prim_paths(self) -> tuple[str, ...]: + """Resolved Newton shape labels in environment-major belt order.""" + return self._require_binding().surface_paths + + @property + def surface_paths(self) -> tuple[str, ...]: + """Alias for :attr:`prim_paths`.""" + return self.prim_paths + + def set_velocities(self, velocities: Any, indices: Any = None) -> None: + """Set signed surface speeds, preserving commands while surfaces are disabled.""" + self._require_binding().set_velocities(velocities, indices) + + def get_velocities(self, indices: Any = None, clone: bool = True) -> wp.array: + """Return effective surface speeds, with disabled surfaces reported as zero.""" + return self._require_binding().get_velocities(indices, clone) + + def get_commanded_velocities(self, indices: Any = None, clone: bool = True) -> wp.array: + """Return staged surface speeds without applying the enabled mask.""" + return self._require_binding().get_commanded_velocities(indices, clone) + + def set_enabled(self, flags: Any, indices: Any = None) -> None: + """Enable or disable selected surfaces without discarding their speed commands.""" + self._require_binding().set_enabled(flags, indices) + + def get_enabled(self, indices: Any = None, clone: bool = True) -> wp.array: + """Return integer enabled flags for selected surfaces.""" + return self._require_binding().get_enabled(indices, clone) + + def get_encoder_positions(self, indices: Any = None, clone: bool = True) -> wp.array: + """Return physics-rate integrated surface travel distances [m].""" + return self._require_binding().get_encoder_positions(indices, clone) + + def reset(self, env_ids: Any = None) -> None: + """Clear stale force and encoder state for selected environments.""" + self._require_binding().reset(env_ids) + + def _request_contact_forces(self, _event: Any) -> None: + """Request solved per-contact forces before the Newton model is finalized.""" + NewtonManager.request_extended_contact_attribute("force") + + def _bind_solver(self, model: Any, contacts: Any) -> None: + """Bind graph callbacks and buffers to the newly initialized solver.""" + previous = self._binding + settings = None + if previous is not None: + settings = ( + previous._command_velocity_host.copy(), + previous._enabled_host.copy(), + ) + previous.close() + self._binding = None + + binding = _SurfaceVelocityBinding(model=model, contacts=contacts, **self._binding_kwargs) + if settings is not None: + velocities, enabled = settings + binding.set_velocities(velocities) + binding.set_enabled(enabled) + self._binding = binding + + def close(self) -> None: + """Deregister lifecycle and graph callbacks; repeated calls are safe.""" + if self._closed: + return + NewtonManager.unregister_solver_init_callback(self._bind_solver) + self._model_init_handle.deregister() + if self._binding is not None: + self._binding.close() + self._binding = None + self._closed = True + + +class _SurfaceVelocityBinding: + """Run one batched moving-surface force pipeline for one Newton model.""" + + def __init__( + self, + model: Any, + contacts: Any, + num_envs: int, + surface_specs: Sequence[SurfaceVelocitySpec], + *, + body_pattern: str, + body_count_per_env: int | None = None, + startup_duration_s: float = 1.0, + env_path_format: str = "/World/envs/env_{}", + ) -> None: + """Initialize the binding before Newton CUDA graph capture. + + Args: + model: Finalized Newton model owned by the active solver. + contacts: Contact buffer owned by the active solver. + num_envs: Number of replicated simulation environments. + surface_specs: Authored surface descriptions in stable within-environment order. + body_pattern: Regular expression selecting bodies that receive surface traction. + body_count_per_env: Expected selected body count per environment, or ``None``. + startup_duration_s: Duration of the initial traction ramp [s]. + env_path_format: Format string resolving one exact environment root from its integer world index. + """ + if not isinstance(num_envs, int) or isinstance(num_envs, bool) or num_envs <= 0: + raise ValueError(f"num_envs must be a positive integer, got {num_envs!r}.") + if not np.isfinite(startup_duration_s) or startup_duration_s <= 0.0: + raise ValueError(f"Conveyor startup duration must be positive, got {startup_duration_s}.") + _validate_env_path_format(env_path_format) + + self._surface_specs = tuple(surface_specs) + _validate_newton_surface_specs(self._surface_specs) + try: + compiled_body_pattern = re.compile(body_pattern) + except re.error as exc: + raise ValueError(f"Invalid body pattern: {body_pattern!r}.") from exc + + if model is None or contacts is None: + raise RuntimeError("The conveyor driver requires an initialized Newton model and contact buffer.") + if contacts.force is None: + raise RuntimeError( + "Newton did not allocate per-contact force reporting. The conveyor driver must request the " + "'force' contact attribute before model finalization." + ) + if model.world_count != num_envs: + raise RuntimeError(f"Newton model has {model.world_count} worlds, expected {num_envs}.") + + self._model = model + self._contacts = contacts + self._device = model.device + self._num_envs = num_envs + self._startup_duration_s = startup_duration_s + self._closed = False + self._validate_backend_buffers() + + surfaces_per_env = len(self._surface_specs) + conveyor_count = num_envs * surfaces_per_env + shape_conveyor = [-1] * model.shape_count + field_type = [0] * conveyor_count + direction = [wp.vec3() for _ in range(conveyor_count)] + pivot_point = [wp.vec3() for _ in range(conveyor_count)] + radius = [1.0] * conveyor_count + surface_normal = [wp.vec3() for _ in range(conveyor_count)] + conveyor_world = [conveyor_id // surfaces_per_env for conveyor_id in range(conveyor_count)] + surface_paths = [""] * conveyor_count + + shape_body = model.shape_body.numpy() + shape_world = model.shape_world.numpy() + shape_transform = model.shape_transform.numpy() + seen_sections: set[tuple[int, int]] = set() + for shape_id, label in enumerate(model.shape_label): + world_id = int(shape_world[shape_id]) + if not 0 <= world_id < num_envs: + continue + matching_specs = [ + index + for index, spec in enumerate(self._surface_specs) + if _shape_belongs_to_prim(label, _resolve_belt_prim_path(spec.prim_path, env_path_format, world_id)) + ] + if not matching_specs: + continue + if len(matching_specs) > 1: + raise RuntimeError(f"Conveyor shape {label!r} matches more than one section specification.") + if int(shape_body[shape_id]) >= 0: + raise ValueError(f"Conveyor shape must be static: {label}") + + spec_id = matching_specs[0] + section_key = (world_id, spec_id) + if section_key in seen_sections: + raise RuntimeError( + f"World {world_id} contains multiple shapes matching conveyor section " + f"{self._surface_specs[spec_id].prim_path!r}." + ) + seen_sections.add(section_key) + + spec = self._surface_specs[spec_id] + conveyor_id = world_id * surfaces_per_env + spec_id + shape_conveyor[shape_id] = conveyor_id + field_type[conveyor_id] = _VELOCITY_FIELD_TYPE_PIVOT if spec.curved else _VELOCITY_FIELD_TYPE_CONSTANT + direction[conveyor_id] = _world_vector(shape_transform[shape_id], spec.direction) + pivot_point[conveyor_id] = _world_point(shape_transform[shape_id], spec.pivot_point) + radius[conveyor_id] = 1.0 if spec.radius is None else spec.radius + surface_normal[conveyor_id] = _world_vector(shape_transform[shape_id], spec.surface_normal) + surface_paths[conveyor_id] = label + + expected_sections = { + (world_id, spec_id) for world_id in range(num_envs) for spec_id in range(len(self._surface_specs)) + } + missing_sections = sorted(expected_sections - seen_sections) + if missing_sections: + details = ", ".join( + f"world {world_id}: {self._surface_specs[spec_id].prim_path}" + for world_id, spec_id in missing_sections[:8] + ) + raise RuntimeError(f"Missing {len(missing_sections)} conveyor collision sections ({details}).") + + body_is_tracked = np.zeros(model.body_count, dtype=np.int32) + tracked_counts = np.zeros(num_envs, dtype=np.int32) + body_world = model.body_world.numpy() + for body_id, label in enumerate(model.body_label): + if compiled_body_pattern.search(label) is None: + continue + world_id = int(body_world[body_id]) + if not 0 <= world_id < num_envs: + raise RuntimeError(f"Transported body {label!r} belongs to invalid world {world_id}.") + body_is_tracked[body_id] = 1 + tracked_counts[world_id] += 1 + + if body_count_per_env is not None: + bad_worlds = np.flatnonzero(tracked_counts != body_count_per_env) + if bad_worlds.size: + details = ", ".join(f"world {world_id}: {tracked_counts[world_id]}" for world_id in bad_worlds[:8]) + raise RuntimeError( + f"Body pattern {body_pattern!r} expected {body_count_per_env} bodies per world ({details})." + ) + if not np.any(body_is_tracked): + raise RuntimeError(f"Body pattern {body_pattern!r} matched no Newton bodies.") + + self._surface_paths = tuple(surface_paths) + self._shape_conveyor = wp.array(shape_conveyor, dtype=wp.int32, device=self._device) + self._body_is_tracked = wp.array(body_is_tracked, dtype=wp.int32, device=self._device) + self._field_type = wp.array(field_type, dtype=wp.int32, device=self._device) + self._direction = wp.array(direction, dtype=wp.vec3, device=self._device) + self._pivot_point = wp.array(pivot_point, dtype=wp.vec3, device=self._device) + self._radius = wp.array(radius, dtype=wp.float32, device=self._device) + self._surface_normal = wp.array(surface_normal, dtype=wp.vec3, device=self._device) + self._conveyor_world = wp.array(conveyor_world, dtype=wp.int32, device=self._device) + + authored_velocity = np.asarray([spec.velocity for spec in self._surface_specs], dtype=np.float32) + authored_enabled = np.asarray([spec.enabled for spec in self._surface_specs], dtype=np.int32) + authored_friction = np.asarray([spec.friction_coefficient for spec in self._surface_specs], dtype=np.float32) + authored_threshold = np.asarray([spec.contact_threshold for spec in self._surface_specs], dtype=np.float32) + self._command_velocity_host = np.tile(authored_velocity, num_envs) + self._enabled_host = np.tile(authored_enabled, num_envs) + self._command_velocity = wp.array(self._command_velocity_host, dtype=wp.float32, device=self._device) + self._enabled = wp.array(self._enabled_host, dtype=wp.int32, device=self._device) + self._effective_velocity = wp.zeros(conveyor_count, dtype=wp.float32, device=self._device) + self._friction = wp.array(np.tile(authored_friction, num_envs), dtype=wp.float32, device=self._device) + self._threshold = wp.array(np.tile(authored_threshold, num_envs), dtype=wp.float32, device=self._device) + self._encoder_position = wp.zeros(conveyor_count, dtype=wp.float32, device=self._device) + self._elapsed_time = wp.zeros(1, dtype=wp.float32, device=self._device) + self._velocity_scale = wp.zeros(1, dtype=wp.float32, device=self._device) + + contact_capacity = contacts.rigid_contact_max + self._contact_force = wp.zeros(contact_capacity, dtype=wp.vec3, device=self._device) + self._belt_contacts = wp.empty(contact_capacity, dtype=BeltContact, device=self._device) + self._body_contact_head = wp.full(model.body_count, -1, dtype=wp.int32, device=self._device) + self._contact_patch_head = wp.full(contact_capacity, -1, dtype=wp.int32, device=self._device) + self._adjusted_normal_force = wp.zeros(contact_capacity, dtype=wp.float32, device=self._device) + self._mass_splitting_scale = wp.zeros(contact_capacity, dtype=wp.float32, device=self._device) + self._body_force = wp.zeros(model.body_count, dtype=wp.spatial_vector, device=self._device) + self._world_mask_host = np.zeros(num_envs, dtype=np.bool_) + self._world_mask = wp.zeros(num_envs, dtype=wp.bool, device=self._device) + self._refresh_effective_velocities() + + NewtonManager.register_state_force_callback(self.apply) + NewtonManager.register_post_solver_substep_callback(self.update) + + @property + def surface_paths(self) -> tuple[str, ...]: + """Resolved Newton shape labels in conveyor-index order.""" + return self._surface_paths + + def set_velocities(self, velocities: Any, indices: Any = None) -> None: + """Set signed surface speeds, preserving commands while surfaces are disabled.""" + selected = self._resolve_indices(indices) + self._command_velocity_host[selected] = self._broadcast_1d(velocities, len(selected), "velocities") + self._command_velocity.assign(self._command_velocity_host) + self._refresh_effective_velocities() + + def get_velocities(self, indices: Any = None, clone: bool = True) -> wp.array: + """Return effective surface speeds, with disabled surfaces reported as zero.""" + return self._get_device_values(self._effective_velocity, indices, clone) + + def get_commanded_velocities(self, indices: Any = None, clone: bool = True) -> wp.array: + """Return staged surface speeds without applying the enabled mask.""" + return self._get_device_values(self._command_velocity, indices, clone) + + def set_enabled(self, flags: Any, indices: Any = None) -> None: + """Enable or disable selected surfaces without discarding their speed commands.""" + selected = self._resolve_indices(indices) + values = self._broadcast_1d(flags, len(selected), "enabled flags") + if not np.all(np.isin(values, (0.0, 1.0))): + raise ValueError(f"Surface enabled flags must contain {len(selected)} boolean values.") + self._enabled_host[selected] = values.astype(np.int32) + self._enabled.assign(self._enabled_host) + self._refresh_effective_velocities() + + def get_enabled(self, indices: Any = None, clone: bool = True) -> wp.array: + """Return integer enabled flags for selected surfaces.""" + return self._get_device_int_values(self._enabled, indices, clone) + + def get_encoder_positions(self, indices: Any = None, clone: bool = True) -> wp.array: + """Return physics-rate integrated surface travel distances [m].""" + return self._get_device_values(self._encoder_position, indices, clone) + + def reset(self, env_ids: Any = None) -> None: + """Clear stale force and encoder state for selected environments. + + A full reset also restarts the global startup ramp. Partial vectorized + resets leave other environments' conveyor forces and startup state intact. + + Args: + env_ids: Environment indices to reset, or ``None`` for every environment. + """ + if env_ids is None: + self._body_force.zero_() + self._encoder_position.zero_() + self._elapsed_time.zero_() + self._velocity_scale.zero_() + return + + ids = _as_numpy(env_ids) + if not np.issubdtype(ids.dtype, np.integer): + raise IndexError(f"Surface reset environment indices must be integers, got {env_ids!r}.") + ids = ids.astype(np.int64, copy=False).reshape(-1) + if np.any((ids < 0) | (ids >= self._num_envs)): + raise IndexError(f"Conveyor reset environment indices are out of range: {ids.tolist()}.") + self._world_mask_host.fill(False) + self._world_mask_host[ids] = True + if np.all(self._world_mask_host): + self.reset() + return + self._world_mask.assign(self._world_mask_host) + wp.launch( + _clear_selected_body_forces, + dim=self._model.body_count, + inputs=[self._model.body_world, self._world_mask], + outputs=[self._body_force], + device=self._device, + ) + wp.launch( + _clear_selected_encoders, + dim=len(self._surface_paths), + inputs=[self._conveyor_world, self._world_mask], + outputs=[self._encoder_position], + device=self._device, + ) + + def close(self) -> None: + """Deregister Newton callbacks and release references held by the driver.""" + if self._closed: + return + NewtonManager.unregister_state_force_callback(self.apply) + NewtonManager.unregister_post_solver_substep_callback(self.update) + self._closed = True + + def apply(self, state) -> None: + """Apply the wrench computed from the preceding physics solve.""" + wp.launch( + _add_body_force, + dim=self._model.body_count, + inputs=[state.body_f, self._body_force], + device=self._device, + ) + + def update(self, solver, contacts, state, dt: float) -> None: + """Read solved contacts and compute the next per-solve conveyor wrench.""" + solver.update_contacts(contacts) + self._body_force.zero_() + self._body_contact_head.fill_(-1) + self._contact_patch_head.fill_(-1) + self._mass_splitting_scale.zero_() + wp.launch( + _advance_startup_scale, + dim=1, + inputs=[dt, self._startup_duration_s], + outputs=[self._elapsed_time, self._velocity_scale], + device=self._device, + ) + wp.launch( + _integrate_encoders, + dim=len(self._surface_paths), + inputs=[dt, self._effective_velocity], + outputs=[self._encoder_position], + device=self._device, + ) + wp.launch( + _extract_linear_force, + dim=self._contacts.rigid_contact_max, + inputs=[contacts.force, self._contact_force], + device=self._device, + ) + wp.launch( + _classify_contacts, + dim=self._contacts.rigid_contact_max, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + contacts.rigid_contact_normal, + contacts.rigid_contact_point0, + contacts.rigid_contact_point1, + self._contact_force, + self._model.shape_body, + self._shape_conveyor, + self._body_is_tracked, + state.body_q, + self._surface_normal, + self._threshold, + ], + outputs=[self._belt_contacts, self._body_contact_head], + device=self._device, + ) + wp.launch( + _prepare_contact_patches, + dim=self._model.body_count, + inputs=[self._belt_contacts, self._body_contact_head, state.body_q, self._model.body_com], + outputs=[self._contact_patch_head, self._adjusted_normal_force, self._mass_splitting_scale], + device=self._device, + ) + wp.launch( + _accumulate_forces, + dim=self._contacts.rigid_contact_max, + inputs=[ + dt, + self._belt_contacts, + state.body_q, + state.body_qd, + self._model.body_com, + self._model.body_inv_mass, + self._model.body_inv_inertia, + self._adjusted_normal_force, + self._mass_splitting_scale, + self._field_type, + self._direction, + self._pivot_point, + self._radius, + self._effective_velocity, + self._friction, + self._velocity_scale, + ], + outputs=[self._body_force], + device=self._device, + ) + + def _validate_backend_buffers(self) -> None: + """Validate every fixed-size Newton buffer consumed by conveyor kernels.""" + model = self._model + contacts = self._contacts + _require_buffer_length("model.shape_body", model.shape_body, model.shape_count) + _require_buffer_length("model.shape_world", model.shape_world, model.shape_count) + _require_buffer_length("model.shape_transform", model.shape_transform, model.shape_count) + _require_buffer_length("model.body_world", model.body_world, model.body_count) + _require_buffer_length("model.body_com", model.body_com, model.body_count) + _require_buffer_length("model.body_inv_mass", model.body_inv_mass, model.body_count) + _require_buffer_length("model.body_inv_inertia", model.body_inv_inertia, model.body_count) + _require_buffer_length("contacts.force", contacts.force, contacts.rigid_contact_max) + _require_buffer_length( + "contacts.rigid_contact_shape0", contacts.rigid_contact_shape0, contacts.rigid_contact_max + ) + _require_buffer_length( + "contacts.rigid_contact_shape1", contacts.rigid_contact_shape1, contacts.rigid_contact_max + ) + _require_buffer_length( + "contacts.rigid_contact_normal", contacts.rigid_contact_normal, contacts.rigid_contact_max + ) + _require_buffer_length( + "contacts.rigid_contact_point0", contacts.rigid_contact_point0, contacts.rigid_contact_max + ) + _require_buffer_length( + "contacts.rigid_contact_point1", contacts.rigid_contact_point1, contacts.rigid_contact_max + ) + _require_buffer_length("contacts.rigid_contact_count", contacts.rigid_contact_count, 1) + + def _resolve_indices(self, indices: Any) -> np.ndarray: + """Normalize and validate a surface index selection.""" + if indices is None: + return np.arange(len(self._surface_paths), dtype=np.int64) + if isinstance(indices, slice): + return np.arange(len(self._surface_paths), dtype=np.int64)[indices] + selected = _as_numpy(indices) + if selected.dtype == np.bool_: + if selected.ndim != 1 or selected.size != len(self._surface_paths): + raise IndexError(f"Boolean surface indices must have length {len(self._surface_paths)}.") + return np.flatnonzero(selected).astype(np.int64) + if not np.issubdtype(selected.dtype, np.integer): + raise IndexError(f"Surface indices must be integers, got {indices!r}.") + selected = selected.astype(np.int64, copy=False).reshape(-1) + if np.any((selected < 0) | (selected >= len(self._surface_paths))): + raise IndexError(f"Conveyor surface indices are out of range: {selected.tolist()}.") + return selected + + def _refresh_effective_velocities(self) -> None: + """Apply the enabled mask at the one device-side command seam.""" + wp.launch( + _update_effective_velocities, + dim=len(self._surface_paths), + inputs=[self._command_velocity, self._enabled], + outputs=[self._effective_velocity], + device=self._device, + ) + + def _get_device_values(self, source: wp.array, indices: Any, clone: bool) -> wp.array: + """Clone a complete device buffer or gather a selected subset.""" + if indices is None: + return wp.clone(source) if clone else source + selected = self._resolve_indices(indices) + selected_device = wp.array(selected, dtype=wp.int32, device=self._device) + values = wp.empty(len(selected), dtype=wp.float32, device=self._device) + if len(selected) > 0: + wp.launch( + _gather_float_values, + dim=len(selected), + inputs=[source, selected_device], + outputs=[values], + device=self._device, + ) + return values + + def _get_device_int_values(self, source: wp.array, indices: Any, clone: bool) -> wp.array: + """Clone an integer device buffer or gather a selected subset.""" + if indices is None: + return wp.clone(source) if clone else source + selected = self._resolve_indices(indices) + selected_device = wp.array(selected, dtype=wp.int32, device=self._device) + values = wp.empty(len(selected), dtype=wp.int32, device=self._device) + if len(selected) > 0: + wp.launch( + _gather_int_values, + dim=len(selected), + inputs=[source, selected_device], + outputs=[values], + device=self._device, + ) + return values + + @staticmethod + def _broadcast_1d(values: Any, count: int, name: str) -> np.ndarray: + """Broadcast one scalar or validate one value per selected surface.""" + array = np.asarray(_as_numpy(values), dtype=np.float32) + if not np.all(np.isfinite(array)): + raise ValueError(f"Conveyor {name} must contain only finite values.") + if array.ndim == 0 or array.size == 1: + return np.full(count, float(array.reshape(-1)[0]), dtype=np.float32) + if array.ndim != 1 or array.size != count: + raise ValueError(f"Conveyor {name} need one value or {count} values, got shape {array.shape}.") + return array 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 a123629772d3..7f106a83ca96 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -267,6 +267,74 @@ def contacts(self): assert NewtonManager._contacts.rigid_contact_max == 2 +def test_in_graph_callback_registration_has_symmetric_cleanup(monkeypatch: pytest.MonkeyPatch) -> None: + """Scene-owned callbacks can deregister safely before Newton manager teardown.""" + + def actuator_callback(): + pass + + def state_force_callback(_state): + pass + + def substep_callback(_solver, _contacts, _state, _dt): + pass + + def solver_init_callback(_model, _contacts): + pass + + monkeypatch.setattr(NewtonManager, "_post_actuator_callbacks", []) + monkeypatch.setattr(NewtonManager, "_state_force_callbacks", []) + monkeypatch.setattr(NewtonManager, "_post_solver_substep_callbacks", []) + monkeypatch.setattr(NewtonManager, "_solver_init_callbacks", []) + + NewtonManager.register_post_actuator_callback(actuator_callback) + NewtonManager.register_state_force_callback(state_force_callback) + NewtonManager.register_post_solver_substep_callback(substep_callback) + NewtonManager.register_solver_init_callback(solver_init_callback) + + assert NewtonManager._post_actuator_callbacks == [actuator_callback] + assert NewtonManager._state_force_callbacks == [state_force_callback] + assert NewtonManager._post_solver_substep_callbacks == [substep_callback] + assert NewtonManager._solver_init_callbacks == [solver_init_callback] + + NewtonManager.unregister_post_actuator_callback(actuator_callback) + NewtonManager.unregister_state_force_callback(state_force_callback) + NewtonManager.unregister_post_solver_substep_callback(substep_callback) + NewtonManager.unregister_solver_init_callback(solver_init_callback) + # Repeated cleanup is intentionally a safe no-op. + NewtonManager.unregister_post_actuator_callback(actuator_callback) + NewtonManager.unregister_state_force_callback(state_force_callback) + NewtonManager.unregister_post_solver_substep_callback(substep_callback) + NewtonManager.unregister_solver_init_callback(solver_init_callback) + + assert NewtonManager._post_actuator_callbacks == [] + assert NewtonManager._state_force_callbacks == [] + assert NewtonManager._post_solver_substep_callbacks == [] + assert NewtonManager._solver_init_callbacks == [] + + +def test_solver_init_callbacks_rebind_to_each_model(monkeypatch: pytest.MonkeyPatch) -> None: + """Solver-init callbacks run once per model, including hard-reset replacements.""" + calls = [] + + def callback(model, contacts): + calls.append((model, contacts)) + + monkeypatch.setattr(NewtonManager, "_solver_init_callbacks", []) + NewtonManager.register_solver_init_callback(callback) + # Duplicate registration must not cause two graph bindings. + NewtonManager.register_solver_init_callback(callback) + + first = (object(), object()) + second = (object(), object()) + for model, contacts in (first, second): + monkeypatch.setattr(NewtonManager, "backend", SimpleNamespace(model=model)) + monkeypatch.setattr(NewtonManager, "_contacts", contacts) + NewtonManager._run_solver_init_callbacks() + + assert calls == [first, second] + + @pytest.mark.parametrize("cloth", [False, True]) @pytest.mark.parametrize("device", test_devices()) def test_queries_share_native_bvhs_and_read_only_through_sdp(monkeypatch, cloth, device): diff --git a/source/isaaclab_newton/test/physics/test_surface_velocity.py b/source/isaaclab_newton/test/physics/test_surface_velocity.py new file mode 100644 index 000000000000..728ddac3df98 --- /dev/null +++ b/source/isaaclab_newton/test/physics/test_surface_velocity.py @@ -0,0 +1,371 @@ +# 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 + +"""Lifecycle tests for Newton surface velocity.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import isaaclab_newton.physics.surface_velocity as surface_module +import numpy as np +import pytest +import warp as wp +from isaaclab_newton.physics.surface_velocity import compute_point_impulse + +from isaaclab.physics import PhysicsEvent, SurfaceVelocitySpec + +_BODY_PATTERN = r"(?:^|/)Cube_?[0-3](?:/|$)" + + +@wp.kernel +def _compute_test_impulses( + target_velocity: wp.array(dtype=wp.vec3), + normal_impulse: wp.array(dtype=wp.float32), + output: wp.array(dtype=wp.vec3), +): + index = wp.tid() + output[index] = compute_point_impulse( + wp.vec3(0.0, 0.0, 1.0), + normal_impulse[index], + wp.vec3(), + target_velocity[index], + 1.0, + wp.mat33(), + wp.vec3(), + 0.5, + 1.0, + ) + + +def _surface_spec(name: str = "Belt") -> SurfaceVelocitySpec: + """Build one valid replicated test belt.""" + return SurfaceVelocitySpec(prim_path=f"{{ENV_REGEX_NS}}/{name}", velocity=0.35, friction_coefficient=0.5) + + +def test_point_impulse_tracks_velocity_and_respects_coulomb_limit() -> None: + """Point traction reaches a small target but clamps large requests to ``mu * normal_impulse``.""" + target_velocity = wp.array([(0.25, 0.0, 0.0), (10.0, 0.0, 0.0)], dtype=wp.vec3, device="cpu") + normal_impulse = wp.array([2.0, 2.0], dtype=wp.float32, device="cpu") + output = wp.zeros(2, dtype=wp.vec3, device="cpu") + + wp.launch( + _compute_test_impulses, + dim=2, + inputs=[target_velocity, normal_impulse], + outputs=[output], + device="cpu", + ) + + np.testing.assert_allclose(output.numpy(), ((0.25, 0.0, 0.0), (1.0, 0.0, 0.0)), atol=1.0e-6) + + +class _FakeCallbackHandle: + def __init__(self) -> None: + self.deregister_count = 0 + + def deregister(self) -> None: + self.deregister_count += 1 + + +class _FakeBinding: + instances = [] + + def __init__(self, model, contacts, **kwargs) -> None: + self.model = model + self.contacts = contacts + self.kwargs = kwargs + self.closed = False + self._command_velocity_host = np.array([0.35, 0.35], dtype=np.float32) + self._enabled_host = np.ones(2, dtype=np.int32) + type(self).instances.append(self) + + def set_velocities(self, values) -> None: + self._command_velocity_host = np.asarray(values, dtype=np.float32).copy() + + def set_enabled(self, values) -> None: + self._enabled_host = np.asarray(values, dtype=np.int32).copy() + + def close(self) -> None: + self.closed = True + + +def test_driver_requests_force_and_rebinds_on_solver_reinitialization(monkeypatch: pytest.MonkeyPatch) -> None: + """The driver binds pre-capture and replaces model-owned buffers after a hard reset.""" + event_callbacks = [] + solver_callbacks = [] + unregistered_solver_callbacks = [] + requested_attributes = [] + callback_handle = _FakeCallbackHandle() + _FakeBinding.instances = [] + + def register_callback(cls, callback, event, order=0, name=None, wrap_weak_ref=True): + event_callbacks.append((callback, event, name)) + return callback_handle + + monkeypatch.setattr(surface_module.NewtonManager, "register_callback", classmethod(register_callback)) + monkeypatch.setattr( + surface_module.NewtonManager, + "register_solver_init_callback", + classmethod(lambda cls, callback: solver_callbacks.append(callback)), + ) + monkeypatch.setattr( + surface_module.NewtonManager, + "unregister_solver_init_callback", + classmethod(lambda cls, callback: unregistered_solver_callbacks.append(callback)), + ) + monkeypatch.setattr( + surface_module.NewtonManager, + "request_extended_contact_attribute", + classmethod(lambda cls, attribute: requested_attributes.append(attribute)), + ) + monkeypatch.setattr(surface_module, "_SurfaceVelocityBinding", _FakeBinding) + + driver = surface_module.SurfaceVelocity(num_envs=2, surface_specs=(_surface_spec(),), body_pattern=_BODY_PATTERN) + + assert driver.specs == (_surface_spec(),) + assert driver.surfaces_per_env == 1 + assert driver.num_surfaces == 2 + assert driver.count == 2 + assert not driver.initialized + + assert [(event, name) for _, event, name in event_callbacks] == [ + (PhysicsEvent.MODEL_INIT, "surface_velocity_contact_attribute") + ] + event_callbacks[0][0](None) + assert requested_attributes == ["force"] + + first_model, first_contacts = object(), object() + solver_callbacks[0](first_model, first_contacts) + first_binding = _FakeBinding.instances[-1] + assert driver.initialized + first_binding.set_velocities([0.2, -0.1]) + first_binding.set_enabled([1, 0]) + + second_model, second_contacts = object(), object() + solver_callbacks[0](second_model, second_contacts) + second_binding = _FakeBinding.instances[-1] + + assert first_binding.closed + assert second_binding.model is second_model + assert second_binding.contacts is second_contacts + np.testing.assert_allclose(second_binding._command_velocity_host, [0.2, -0.1]) + np.testing.assert_array_equal(second_binding._enabled_host, [1, 0]) + + driver.close() + driver.close() + assert second_binding.closed + assert unregistered_solver_callbacks == [solver_callbacks[0]] + assert callback_handle.deregister_count == 1 + + +def test_unbound_driver_rejects_control_calls(monkeypatch: pytest.MonkeyPatch) -> None: + """Control methods are unavailable until the solver-init callback creates a binding.""" + callback_handle = _FakeCallbackHandle() + monkeypatch.setattr( + surface_module.NewtonManager, + "register_callback", + classmethod(lambda cls, *args, **kwargs: callback_handle), + ) + monkeypatch.setattr( + surface_module.NewtonManager, + "register_solver_init_callback", + classmethod(lambda cls, callback: None), + ) + monkeypatch.setattr( + surface_module.NewtonManager, + "unregister_solver_init_callback", + classmethod(lambda cls, callback: None), + ) + + driver = surface_module.SurfaceVelocity(num_envs=1, surface_specs=(_surface_spec(),), body_pattern=_BODY_PATTERN) + with pytest.raises(RuntimeError, match="not bound"): + driver.set_velocities(0.2) + driver.close() + + +def test_driver_rejects_invalid_specs_before_registering_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + """Invalid descriptions cannot leave lifecycle callbacks behind.""" + registered = [] + monkeypatch.setattr( + surface_module.NewtonManager, + "register_callback", + classmethod(lambda cls, *args, **kwargs: registered.append(args)), + ) + + with pytest.raises(ValueError, match="At least one"): + surface_module.SurfaceVelocity(num_envs=1, surface_specs=(), body_pattern=_BODY_PATTERN) + with pytest.raises(TypeError, match="SurfaceVelocitySpec"): + surface_module.SurfaceVelocity(num_envs=1, surface_specs=(object(),), body_pattern=_BODY_PATTERN) + with pytest.raises(ValueError, match="unique"): + surface_module.SurfaceVelocity( + num_envs=1, surface_specs=(_surface_spec(), _surface_spec()), body_pattern=_BODY_PATTERN + ) + with pytest.raises(ValueError, match="ancestors"): + surface_module.SurfaceVelocity( + num_envs=1, + surface_specs=( + SurfaceVelocitySpec(prim_path="{ENV_REGEX_NS}/Belt"), + SurfaceVelocitySpec(prim_path="{ENV_REGEX_NS}/Belt/Child"), + ), + body_pattern=_BODY_PATTERN, + ) + with pytest.raises(ValueError, match="explicit positive radius"): + surface_module.SurfaceVelocity( + num_envs=1, + surface_specs=(SurfaceVelocitySpec(prim_path="{ENV_REGEX_NS}/Curve", curved=True),), + body_pattern=_BODY_PATTERN, + ) + with pytest.raises(ValueError, match="Replicated conveyor environments"): + surface_module.SurfaceVelocity( + num_envs=2, + surface_specs=(SurfaceVelocitySpec(prim_path="/World/Shared/Belt"),), + body_pattern=_BODY_PATTERN, + ) + with pytest.raises(ValueError, match="env_path_format"): + surface_module.SurfaceVelocity( + num_envs=1, + surface_specs=(_surface_spec(),), + body_pattern=_BODY_PATTERN, + env_path_format="/World/envs/env_.*", + ) + + assert registered == [] + + +def test_belt_paths_are_exact_and_environment_scoped() -> None: + """A descriptor cannot bind a same-named shape outside the replicated environment root.""" + resolve = surface_module._resolve_belt_prim_path + + assert resolve("{ENV_REGEX_NS}/Belt", "/World/envs/env_{}", 0) == "/World/envs/env_0/Belt" + assert resolve("{ENV_REGEX_NS}/Nested/Belt", "/World/envs/env_{}", 123) == ("/World/envs/env_123/Nested/Belt") + assert resolve("/World/Shared/Belt", "/World/envs/env_{}", 7) == "/World/Shared/Belt" + belongs = surface_module._shape_belongs_to_prim + assert belongs("/World/envs/env_0/Belt/geometry/mesh", "/World/envs/env_0/Belt") + assert not belongs("/World/props/Belt/geometry/mesh", "/World/envs/env_0/Belt") + assert not belongs("/World/envs/env_0/Nested/Belt", "/World/envs/env_0/Belt") + + +def test_driver_cleans_up_model_callback_when_solver_registration_fails(monkeypatch: pytest.MonkeyPatch) -> None: + """A lifecycle registration failure cannot leave a partially active driver.""" + callback_handle = _FakeCallbackHandle() + monkeypatch.setattr( + surface_module.NewtonManager, + "register_callback", + classmethod(lambda cls, *args, **kwargs: callback_handle), + ) + + def fail_registration(cls, callback): + raise RuntimeError("solver callback unavailable") + + monkeypatch.setattr( + surface_module.NewtonManager, + "register_solver_init_callback", + classmethod(fail_registration), + ) + + with pytest.raises(RuntimeError, match="solver callback unavailable"): + surface_module.SurfaceVelocity(num_envs=1, surface_specs=(_surface_spec(),), body_pattern=_BODY_PATTERN) + + assert callback_handle.deregister_count == 1 + + +def test_binding_uses_deterministic_environment_major_belt_indices(monkeypatch: pytest.MonkeyPatch) -> None: + """Newton discovery order cannot reorder commands or encoder rows after a rebuild.""" + monkeypatch.setattr( + surface_module.NewtonManager, + "register_state_force_callback", + classmethod(lambda cls, callback: None), + ) + monkeypatch.setattr( + surface_module.NewtonManager, + "register_post_solver_substep_callback", + classmethod(lambda cls, callback: None), + ) + monkeypatch.setattr( + surface_module.NewtonManager, + "unregister_state_force_callback", + classmethod(lambda cls, callback: None), + ) + monkeypatch.setattr( + surface_module.NewtonManager, + "unregister_post_solver_substep_callback", + classmethod(lambda cls, callback: None), + ) + + shape_labels = ( + "/World/envs/env_1/BeltB", + "/World/envs/env_0/BeltA", + "/World/envs/env_1/BeltA", + "/World/envs/env_0/BeltB", + ) + shape_count = len(shape_labels) + identity = wp.transform(wp.vec3(), wp.quat_identity()) + model = SimpleNamespace( + world_count=2, + device="cpu", + shape_count=shape_count, + shape_label=shape_labels, + shape_body=wp.full(shape_count, -1, dtype=wp.int32, device="cpu"), + shape_world=wp.array([1, 0, 1, 0], dtype=wp.int32, device="cpu"), + shape_transform=wp.array([identity] * shape_count, dtype=wp.transform, device="cpu"), + body_count=2, + body_label=("/World/envs/env_0/Cube0", "/World/envs/env_1/Cube0"), + body_world=wp.array([0, 1], dtype=wp.int32, device="cpu"), + body_com=wp.zeros(2, dtype=wp.vec3, device="cpu"), + body_inv_mass=wp.ones(2, dtype=wp.float32, device="cpu"), + body_inv_inertia=wp.array([wp.mat33(1.0)] * 2, dtype=wp.mat33, device="cpu"), + ) + contact_capacity = 4 + contacts = SimpleNamespace( + rigid_contact_max=contact_capacity, + force=wp.zeros(contact_capacity, dtype=wp.spatial_vector, device="cpu"), + rigid_contact_shape0=wp.full(contact_capacity, -1, dtype=wp.int32, device="cpu"), + rigid_contact_shape1=wp.full(contact_capacity, -1, dtype=wp.int32, device="cpu"), + rigid_contact_normal=wp.zeros(contact_capacity, dtype=wp.vec3, device="cpu"), + rigid_contact_point0=wp.zeros(contact_capacity, dtype=wp.vec3, device="cpu"), + rigid_contact_point1=wp.zeros(contact_capacity, dtype=wp.vec3, device="cpu"), + rigid_contact_count=wp.zeros(1, dtype=wp.int32, device="cpu"), + ) + specs = ( + SurfaceVelocitySpec( + prim_path="{ENV_REGEX_NS}/BeltA", + velocity=0.1, + friction_coefficient=0.4, + contact_threshold=0.98, + ), + SurfaceVelocitySpec( + prim_path="{ENV_REGEX_NS}/BeltB", + velocity=-0.2, + enabled=False, + friction_coefficient=0.6, + contact_threshold=0.99, + ), + ) + + binding = surface_module._SurfaceVelocityBinding( + model=model, + contacts=contacts, + num_envs=2, + surface_specs=specs, + body_pattern=_BODY_PATTERN, + body_count_per_env=1, + ) + try: + assert binding.surface_paths == ( + "/World/envs/env_0/BeltA", + "/World/envs/env_0/BeltB", + "/World/envs/env_1/BeltA", + "/World/envs/env_1/BeltB", + ) + np.testing.assert_array_equal(binding._shape_conveyor.numpy(), [3, 0, 2, 1]) + np.testing.assert_array_equal(binding._conveyor_world.numpy(), [0, 0, 1, 1]) + np.testing.assert_allclose(binding._command_velocity_host, [0.1, -0.2, 0.1, -0.2]) + np.testing.assert_array_equal(binding._enabled_host, [1, 0, 1, 0]) + np.testing.assert_allclose(binding._friction.numpy(), [0.4, 0.6, 0.4, 0.6]) + np.testing.assert_allclose(binding._threshold.numpy(), [0.98, 0.99, 0.98, 0.99]) + np.testing.assert_array_equal(binding.get_enabled(indices=[3, 0]).numpy(), [0, 1]) + finally: + binding.close() diff --git a/source/isaaclab_physx/changelog.d/maximiliank-surface-velocity.minor.rst b/source/isaaclab_physx/changelog.d/maximiliank-surface-velocity.minor.rst new file mode 100644 index 000000000000..33debf44ed07 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/maximiliank-surface-velocity.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added native PhysX surface-velocity schema authoring and runtime control under + ``isaaclab_physx.physics.surface_velocity``. diff --git a/source/isaaclab_physx/isaaclab_physx/physics/__init__.pyi b/source/isaaclab_physx/isaaclab_physx/physics/__init__.pyi index abdeb42beedb..2b9c55d6d392 100644 --- a/source/isaaclab_physx/isaaclab_physx/physics/__init__.pyi +++ b/source/isaaclab_physx/isaaclab_physx/physics/__init__.pyi @@ -8,7 +8,19 @@ __all__ = [ "IsaacEvents", "PhysxCfg", "PhysxBackendCfg", + "PhysxSurfaceVelocityTwist", + "SurfaceVelocity", + "apply_surface_velocity_api", + "compute_surface_velocity_twist", + "resolve_surface_velocity_paths", ] from .physx_manager import PhysxManager, IsaacEvents from .physx_manager_cfg import PhysxCfg, PhysxBackendCfg +from .surface_velocity import ( + PhysxSurfaceVelocityTwist, + SurfaceVelocity, + apply_surface_velocity_api, + compute_surface_velocity_twist, + resolve_surface_velocity_paths, +) diff --git a/source/isaaclab_physx/isaaclab_physx/physics/surface_velocity.py b/source/isaaclab_physx/isaaclab_physx/physics/surface_velocity.py new file mode 100644 index 000000000000..68b382594ae1 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/physics/surface_velocity.py @@ -0,0 +1,532 @@ +# 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 + +"""CPU-only native PhysX surface-velocity control. + +The module deliberately keeps PhysX schema imports behind authoring and binding calls. Its +geometry conversion and host-side control state therefore remain usable in import-light tests +and tooling that do not launch Kit. + +The task configuration rejects GPU dynamics: the supported native ``PhysxSurfaceVelocityAPI`` +path can lose conveyor contacts there. The force-driven Newton adapter is the scalable GPU path. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Protocol + +import numpy as np + +from isaaclab.physics import SurfaceVelocitySpec + +_ENV_REGEX_NS = "{ENV_REGEX_NS}" + + +@dataclass(frozen=True, slots=True) +class PhysxSurfaceVelocityTwist: + """Local PhysX surface twist. + + Attributes: + linear_velocity: Local linear surface velocity [m/s]. + angular_velocity_deg: Local angular surface velocity [deg/s]. + """ + + linear_velocity: tuple[float, float, float] + angular_velocity_deg: tuple[float, float, float] + + +def compute_surface_velocity_twist( + spec: SurfaceVelocitySpec, velocity: float | None = None +) -> PhysxSurfaceVelocityTwist: + """Convert one backend-neutral belt command to a local PhysX surface twist. + + Straight belts map their signed speed to a normalized linear direction. Curved belts map + speed over radius to PhysX's degree-per-second angular convention. PhysX rotates the surface + field about the rigid-body origin, so ``-omega x pivot`` is included in the linear component + to move the instantaneous center to :attr:`SurfaceVelocitySpec.pivot_point`. + + Args: + spec: Authored conveyor intent in the collision prim's local frame. + velocity: Optional signed speed override [m/s]. Defaults to ``spec.velocity``. + + Returns: + Local linear and angular PhysX surface velocities. + + Raises: + ValueError: If the speed is not finite, the direction is zero, or a curved belt has no + positive radius. + """ + speed = spec.velocity if velocity is None else _finite_float("velocity", velocity) + direction = np.asarray(spec.direction, dtype=np.float64) + magnitude = float(np.linalg.norm(direction)) + if not math.isfinite(magnitude) or magnitude <= 1.0e-8: + raise ValueError(f"Conveyor direction must be finite and non-zero, got {spec.direction!r}.") + unit_direction = direction / magnitude + + if not spec.curved: + linear = unit_direction * speed + return PhysxSurfaceVelocityTwist(tuple(float(value) for value in linear), (0.0, 0.0, 0.0)) + + if spec.radius is None or not math.isfinite(spec.radius) or spec.radius <= 0.0: + raise ValueError("A curved PhysX conveyor requires a finite positive radius.") + angular_rad = unit_direction * (speed / spec.radius) + pivot = np.asarray(spec.pivot_point, dtype=np.float64) + linear = -np.cross(angular_rad, pivot) + angular_deg = np.degrees(angular_rad) + return PhysxSurfaceVelocityTwist( + tuple(float(value) for value in linear), tuple(float(value) for value in angular_deg) + ) + + +def apply_surface_velocity_api( + prim_or_path: Any, + spec: SurfaceVelocitySpec, + *, + velocity_scale: float = 0.0, + stage: Any | None = None, +) -> None: + """Apply and author a kinematic PhysX surface-velocity API on one belt prim. + + This function is intended for a task spawner, before PhysX parses the stage. It lazily imports + ``pxr.PhysxSchema`` and applies both the rigid-body and surface-velocity schemas. The rigid body + is made kinematic, surface velocities are authored in local space, and the initial command is + scaled explicitly. A zero default prevents pre-runtime motion during simulation warmup. + + Args: + prim_or_path: USD prim object or exact prim path. + spec: Conveyor intent used to author the local twist. + velocity_scale: Finite multiplier for the initial signed surface speed. + stage: Optional USD stage used when ``prim_or_path`` is a string. Defaults to the current + Isaac Lab stage. + + Raises: + RuntimeError: If the PhysX schema is unavailable or the target prim does not exist. + ValueError: If ``velocity_scale`` is not finite. + """ + scale = _finite_float("velocity_scale", velocity_scale) + twist = compute_surface_velocity_twist(spec, velocity=spec.velocity * scale) + binding = _PhysxSchemaSurfaceWriter((prim_or_path,), stage=stage, apply_api=True) + try: + binding.write(0, enabled=spec.enabled, twist=twist) + finally: + binding.close() + + +def resolve_surface_velocity_paths( + num_envs: int, + surface_specs: Sequence[SurfaceVelocitySpec], + env_path_format: str = "/World/envs/env_{}", +) -> tuple[str, ...]: + """Resolve conveyor templates to exact environment-major prim paths. + + Args: + num_envs: Number of replicated environments. + surface_specs: Within-environment surface descriptions. + env_path_format: Exact environment path format containing one ``{}`` field. + + Returns: + Exact paths ordered by environment, then by ``surface_specs`` order. + + Raises: + ValueError: If inputs cannot produce one unique path per environment and belt. + """ + if not isinstance(num_envs, int) or isinstance(num_envs, bool) or num_envs <= 0: + raise ValueError(f"Conveyor num_envs must be a positive integer, got {num_envs!r}.") + specs = tuple(surface_specs) + if not specs or not all(isinstance(spec, SurfaceVelocitySpec) for spec in specs): + raise ValueError("surface_specs must contain at least one SurfaceVelocitySpec.") + if not isinstance(env_path_format, str) or env_path_format.count("{}") != 1: + raise ValueError(f"Conveyor env_path_format must contain exactly one '{{}}', got {env_path_format!r}.") + try: + env_paths = tuple(env_path_format.format(env_id) for env_id in range(num_envs)) + except (IndexError, KeyError, ValueError) as exc: + raise ValueError(f"Invalid conveyor env_path_format: {env_path_format!r}.") from exc + if any(not path.startswith("/") or "{" in path or "}" in path for path in env_paths): + raise ValueError(f"Conveyor env_path_format must produce exact absolute paths, got {env_path_format!r}.") + + templates = tuple(spec.prim_path for spec in specs) + if len(set(templates)) != len(templates): + raise ValueError("Conveyor belt spec paths must be unique within an environment.") + if num_envs > 1 and any(_ENV_REGEX_NS not in path for path in templates): + raise ValueError("Replicated PhysX conveyors require every belt prim_path to use {ENV_REGEX_NS}.") + + paths = tuple(template.replace(_ENV_REGEX_NS, env_path) for env_path in env_paths for template in templates) + if len(set(paths)) != len(paths): + raise ValueError("Resolved PhysX conveyor prim paths must be unique.") + return paths + + +class SurfaceVelocity: + """Host-side CPU reference facade for native PhysX surface velocity. + + The facade binds exact environment-major paths whose schemas were already authored by + :func:`apply_surface_velocity_api`. Call :meth:`start` to register physics-rate updates, + or call :meth:`update` manually. Commands and enabled state survive full resets; a full reset + clears encoders and restarts the one-second startup ramp. This facade authors USD attributes + on the host and is not a GPU conveyor implementation. + """ + + def __init__( + self, + num_envs: int, + surface_specs: Sequence[SurfaceVelocitySpec], + *, + env_path_format: str = "/World/envs/env_{}", + startup_duration_s: float = 1.0, + stage: Any | None = None, + writer: Any | None = None, + ) -> None: + """Bind native surface attributes and initialize host-side control state. + + Args: + num_envs: Number of replicated environments. + surface_specs: Within-environment surface descriptions. + env_path_format: Exact replicated environment path format. + startup_duration_s: Duration of the global surface-speed ramp [s]. + stage: Optional USD stage used by the default schema writer. + writer: Optional writer implementing ``write(index, enabled=..., twist=...)`` and + ``close()``. This import-light seam is primarily for focused tests. + """ + duration = _finite_float("startup_duration_s", startup_duration_s) + if duration <= 0.0: + raise ValueError(f"Conveyor startup_duration_s must be positive, got {startup_duration_s!r}.") + self._num_envs = num_envs + self._surface_specs = tuple(surface_specs) + self._surface_paths = resolve_surface_velocity_paths(num_envs, self._surface_specs, env_path_format) + self._surfaces_per_env = len(self._surface_specs) + self._startup_duration_s = duration + self._elapsed_time = 0.0 + self._velocity_scale = 0.0 + self._closed = False + self._callback_handle: Any | None = None + self._writer: _SurfaceVelocityWriter = ( + writer + if writer is not None + else _PhysxSchemaSurfaceWriter(self._surface_paths, stage=stage, apply_api=False) + ) + + self._command_velocity = np.tile( + np.asarray([spec.velocity for spec in self._surface_specs], dtype=np.float32), self._num_envs + ) + self._enabled = np.tile( + np.asarray([spec.enabled for spec in self._surface_specs], dtype=np.bool_), self._num_envs + ) + self._encoder_position = np.zeros(self.num_surfaces, dtype=np.float32) + self._last_authored: list[tuple[bool, PhysxSurfaceVelocityTwist] | None] = [None] * self.num_surfaces + self._flush(force=True) + + @property + def specs(self) -> tuple[SurfaceVelocitySpec, ...]: + """Return authored descriptions in stable within-environment order.""" + return self._surface_specs + + @property + def surfaces_per_env(self) -> int: + """Return the number of authored surfaces per environment.""" + return self._surfaces_per_env + + @property + def prim_paths(self) -> tuple[str, ...]: + """Return exact PhysX prim paths in environment-major order.""" + return self._surface_paths + + @property + def surface_paths(self) -> tuple[str, ...]: + """Return an alias for :attr:`prim_paths`.""" + return self.prim_paths + + @property + def num_surfaces(self) -> int: + """Return the total number of bound conveyor surfaces.""" + return len(self._surface_paths) + + @property + def count(self) -> int: + """Return an alias for :attr:`num_surfaces`.""" + return self.num_surfaces + + @property + def initialized(self) -> bool: + """Return whether the facade still owns live schema bindings.""" + return not self._closed + + def start(self) -> None: + """Register one lazy PhysX post-step callback; repeated calls are safe.""" + self._require_open() + if self._callback_handle is not None: + return + from .physx_manager import IsaacEvents, PhysxManager + + self._callback_handle = PhysxManager.register_callback( + self.update, + IsaacEvents.POST_PHYSICS_STEP, + name="physx_surface_velocity", + ) + + def update(self, dt: float) -> None: + """Advance encoder state and the startup ramp by one physics step. + + Args: + dt: Positive physics step duration [s]. + """ + self._require_open() + step = _finite_float("physics dt", dt) + if step <= 0.0: + raise ValueError(f"Conveyor physics dt must be positive, got {dt!r}.") + effective_velocity = self._command_velocity * self._enabled + self._encoder_position += np.asarray(step * effective_velocity, dtype=np.float32) + self._elapsed_time = min(self._startup_duration_s, self._elapsed_time + step) + self._velocity_scale = self._elapsed_time / self._startup_duration_s + self._flush() + + def set_velocities(self, velocities: Any, indices: Any = None) -> None: + """Set signed speeds [m/s], preserving commands while belts are disabled.""" + selected = self._resolve_indices(indices) + self._command_velocity[selected] = self._broadcast_finite(velocities, len(selected), "velocities") + self._flush(selected) + + def get_velocities(self, indices: Any = None, clone: bool = True) -> np.ndarray: + """Return effective speeds with disabled belts reported as zero.""" + return self._get_values(self._command_velocity * self._enabled, indices, clone) + + def get_commanded_velocities(self, indices: Any = None, clone: bool = True) -> np.ndarray: + """Return commanded speeds before applying the enabled mask.""" + return self._get_values(self._command_velocity, indices, clone) + + def set_enabled(self, flags: Any, indices: Any = None) -> None: + """Enable or disable belts without discarding their commands.""" + selected = self._resolve_indices(indices) + values = _as_numpy(flags) + if values.ndim == 0: + values = np.full(len(selected), values.item()) + else: + values = values.reshape(-1) + if values.size == 1: + values = np.full(len(selected), values.item()) + if values.size != len(selected) or not np.all(np.isin(values, (False, True, 0, 1))): + raise ValueError(f"Conveyor enabled flags must contain {len(selected)} boolean values.") + self._enabled[selected] = values.astype(np.bool_) + self._flush(selected) + + def get_enabled(self, indices: Any = None, clone: bool = True) -> np.ndarray: + """Return enabled flags as integer values.""" + return self._get_values(self._enabled.astype(np.int32), indices, clone) + + def get_encoder_positions(self, indices: Any = None, clone: bool = True) -> np.ndarray: + """Return physics-rate integrated commanded belt travel [m].""" + return self._get_values(self._encoder_position, indices, clone) + + def reset(self, env_ids: Any = None) -> None: + """Clear selected encoders and restart the global ramp on a full reset. + + Velocity commands and enabled state are deliberately preserved, including across a full + hard-reset path. Partial resets do not disturb the global ramp used by other environments. + + Args: + env_ids: Environment indices to reset, or ``None`` for all environments. + """ + self._require_open() + ids = self._resolve_env_ids(env_ids) + rows = (ids[:, None] * self._surfaces_per_env + np.arange(self._surfaces_per_env)[None, :]).reshape(-1) + self._encoder_position[rows] = 0.0 + if len(np.unique(ids)) == self._num_envs: + self._elapsed_time = 0.0 + self._velocity_scale = 0.0 + self._flush(force=True) + + def close(self) -> None: + """Deregister callbacks, disable authored motion, and release bindings safely.""" + if self._closed: + return + if self._callback_handle is not None: + self._callback_handle.deregister() + self._callback_handle = None + zero = PhysxSurfaceVelocityTwist((0.0, 0.0, 0.0), (0.0, 0.0, 0.0)) + for index in range(self.num_surfaces): + self._writer.write(index, enabled=False, twist=zero) + self._writer.close() + self._closed = True + + def _flush(self, indices: Any = None, *, force: bool = False) -> None: + """Author scaled effective commands for selected rows.""" + selected = self._resolve_indices(indices) + for index in selected: + enabled = bool(self._enabled[index]) + spec = self._surface_specs[index % self._surfaces_per_env] + speed = float(self._command_velocity[index]) * self._velocity_scale if enabled else 0.0 + twist = compute_surface_velocity_twist(spec, velocity=speed) + state = (enabled, twist) + if force or state != self._last_authored[index]: + self._writer.write(int(index), enabled=enabled, twist=twist) + self._last_authored[index] = state + + def _resolve_indices(self, indices: Any) -> np.ndarray: + """Normalize and validate a belt row selection.""" + self._require_open() + if indices is None: + return np.arange(self.num_surfaces, dtype=np.int64) + if isinstance(indices, slice): + return np.arange(self.num_surfaces, dtype=np.int64)[indices] + selected = _as_numpy(indices) + if selected.dtype == np.bool_: + if selected.ndim != 1 or selected.size != self.num_surfaces: + raise IndexError(f"Boolean surface indices must have length {self.num_surfaces}.") + return np.flatnonzero(selected).astype(np.int64) + if not np.issubdtype(selected.dtype, np.integer): + raise IndexError(f"Conveyor indices must be integers, got {indices!r}.") + selected = selected.astype(np.int64, copy=False).reshape(-1) + if np.any((selected < 0) | (selected >= self.num_surfaces)): + raise IndexError(f"Conveyor surface indices are out of range: {selected.tolist()}.") + return selected + + def _resolve_env_ids(self, env_ids: Any) -> np.ndarray: + """Normalize and validate an environment selection.""" + if env_ids is None: + return np.arange(self._num_envs, dtype=np.int64) + ids = _as_numpy(env_ids) + if not np.issubdtype(ids.dtype, np.integer): + raise IndexError(f"Conveyor reset environment indices must be integers, got {env_ids!r}.") + ids = ids.astype(np.int64, copy=False).reshape(-1) + if np.any((ids < 0) | (ids >= self._num_envs)): + raise IndexError(f"Conveyor reset environment indices are out of range: {ids.tolist()}.") + return ids + + @staticmethod + def _broadcast_finite(values: Any, count: int, name: str) -> np.ndarray: + """Return a finite float32 vector with scalar broadcasting.""" + try: + result = np.asarray(_as_numpy(values), dtype=np.float32) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"Conveyor {name} must contain finite numeric values.") from exc + if result.ndim == 0: + result = np.full(count, result.item(), dtype=np.float32) + else: + result = result.reshape(-1) + if result.size == 1: + result = np.full(count, result.item(), dtype=np.float32) + if result.size != count or not np.all(np.isfinite(result)): + raise ValueError(f"Conveyor {name} must contain one or {count} finite values.") + return result + + def _get_values(self, source: np.ndarray, indices: Any, clone: bool) -> np.ndarray: + """Return a complete buffer or a selected copy.""" + self._require_open() + if indices is None: + return source.copy() if clone else source + return source[self._resolve_indices(indices)].copy() + + def _require_open(self) -> None: + """Fail predictably after backend resources have been released.""" + if self._closed: + raise RuntimeError("The PhysX conveyor surface facade is closed.") + + +class _SurfaceVelocityWriter(Protocol): + """Minimal schema-writer seam used by the runtime facade.""" + + def write(self, index: int, *, enabled: bool, twist: PhysxSurfaceVelocityTwist) -> None: + """Author one bound surface state.""" + ... + + def close(self) -> None: + """Release retained schema attributes.""" + ... + + +class _PhysxSchemaSurfaceWriter: + """Cached USD attribute writer with all schema imports kept lazy.""" + + def __init__(self, prims_or_paths: Sequence[Any], *, stage: Any | None, apply_api: bool) -> None: + """Resolve prims, optionally apply schemas, and cache surface attributes.""" + try: + from pxr import Gf, PhysxSchema, UsdPhysics + except ImportError as exc: + raise RuntimeError( + "PhysxSurfaceVelocityAPI is unavailable. Use this backend inside an Isaac Sim process with PhysX " + "schemas." + ) from exc + + if stage is None and any(isinstance(value, str) for value in prims_or_paths): + from isaaclab.sim.utils.stage import get_current_stage + + stage = get_current_stage() + self._gf = Gf + self._attributes: list[tuple[Any, Any, Any]] = [] + for value in prims_or_paths: + prim = stage.GetPrimAtPath(value) if isinstance(value, str) else value + if prim is None or not prim.IsValid(): + raise RuntimeError(f"Cannot bind PhysX conveyor surface: prim {value!r} does not exist.") + + rigid_body = UsdPhysics.RigidBodyAPI(prim) + if apply_api and not prim.HasAPI(UsdPhysics.RigidBodyAPI): + rigid_body = UsdPhysics.RigidBodyAPI.Apply(prim) + elif not prim.HasAPI(UsdPhysics.RigidBodyAPI): + raise RuntimeError(f"PhysX conveyor prim {prim.GetPath()} has no authored RigidBodyAPI.") + if apply_api: + rigid_body.CreateRigidBodyEnabledAttr().Set(True) + rigid_body.CreateKinematicEnabledAttr().Set(True) + elif rigid_body.GetKinematicEnabledAttr().Get() is not True: + raise RuntimeError(f"PhysX conveyor prim {prim.GetPath()} must be an authored kinematic rigid body.") + + if prim.HasAPI(PhysxSchema.PhysxSurfaceVelocityAPI): + surface_api = PhysxSchema.PhysxSurfaceVelocityAPI(prim) + elif apply_api: + surface_api = PhysxSchema.PhysxSurfaceVelocityAPI.Apply(prim) + else: + raise RuntimeError( + f"PhysX conveyor prim {prim.GetPath()} has no authored PhysxSurfaceVelocityAPI; " + "call apply_surface_velocity_api from its spawner before simulation starts." + ) + surface_api.CreateSurfaceVelocityLocalSpaceAttr().Set(True) + self._attributes.append( + ( + surface_api.CreateSurfaceVelocityEnabledAttr(), + surface_api.CreateSurfaceVelocityAttr(), + surface_api.CreateSurfaceAngularVelocityAttr(), + ) + ) + + def write(self, index: int, *, enabled: bool, twist: PhysxSurfaceVelocityTwist) -> None: + """Author one cached local-space surface twist.""" + enabled_attr, linear_attr, angular_attr = self._attributes[index] + # PhysX caches the contact-modification flag when a surface-velocity + # shape is parsed. Cycle it around command changes, matching the + # official Isaac Conveyor node, so live updates are picked up without + # leaving a stale contact modifier attached to the shape. + enabled_attr.Set(False) + linear_attr.Set(self._gf.Vec3f(*twist.linear_velocity)) + angular_attr.Set(self._gf.Vec3f(*twist.angular_velocity_deg)) + enabled_attr.Set(enabled) + + def close(self) -> None: + """Release cached schema attribute handles.""" + self._attributes.clear() + + +def _finite_float(name: str, value: Any) -> float: + """Coerce one finite scalar with a stable validation error.""" + try: + result = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"Conveyor {name} must be finite, got {value!r}.") from exc + if not math.isfinite(result): + raise ValueError(f"Conveyor {name} must be finite, got {value!r}.") + return result + + +def _as_numpy(values: Any) -> np.ndarray: + """Move supported tensor-like values to a host NumPy array without importing their libraries.""" + if isinstance(values, np.ndarray): + return values + if hasattr(values, "detach"): + values = values.detach() + if hasattr(values, "cpu"): + values = values.cpu() + if hasattr(values, "numpy"): + return np.asarray(values.numpy()) + return np.asarray(values) diff --git a/source/isaaclab_physx/test/sim/test_surface_velocity.py b/source/isaaclab_physx/test/sim/test_surface_velocity.py new file mode 100644 index 000000000000..1bde98d6f559 --- /dev/null +++ b/source/isaaclab_physx/test/sim/test_surface_velocity.py @@ -0,0 +1,229 @@ +# 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 + +"""Import-light tests for native PhysX surface velocity.""" + +from __future__ import annotations + +import math +import sys +from types import ModuleType + +import isaaclab_physx.physics.surface_velocity as surface_module +import numpy as np +import pytest +import torch + +from isaaclab.physics import SurfaceVelocitySpec + + +class _FakeWriter: + """Record authored states without importing USD or PhysX schemas.""" + + def __init__(self) -> None: + self.writes: list[tuple[int, bool, surface_module.PhysxSurfaceVelocityTwist]] = [] + self.close_count = 0 + + def write(self, index: int, *, enabled: bool, twist: surface_module.PhysxSurfaceVelocityTwist) -> None: + self.writes.append((index, enabled, twist)) + + def close(self) -> None: + self.close_count += 1 + + +def _surface_spec(name: str = "Belt", **kwargs) -> SurfaceVelocitySpec: + """Return one replicated belt description for facade tests.""" + return SurfaceVelocitySpec(prim_path=f"{{ENV_REGEX_NS}}/{name}", velocity=0.4, **kwargs) + + +def test_twist_conversion_normalizes_straight_direction() -> None: + """Straight surface speed is independent of the authored direction magnitude.""" + spec = _surface_spec(direction=(3.0, 4.0, 0.0)) + + twist = surface_module.compute_surface_velocity_twist(spec, velocity=2.0) + + np.testing.assert_allclose(twist.linear_velocity, (1.2, 1.6, 0.0)) + assert twist.angular_velocity_deg == (0.0, 0.0, 0.0) + + +def test_twist_conversion_uses_degrees_and_compensates_curved_pivot() -> None: + """Curved belts rotate about their local pivot rather than the rigid-body origin.""" + spec = _surface_spec( + direction=(0.0, 0.0, 2.0), + curved=True, + radius=2.0, + pivot_point=(2.0, 0.0, 0.0), + ) + + twist = surface_module.compute_surface_velocity_twist(spec, velocity=math.pi) + + np.testing.assert_allclose(twist.angular_velocity_deg, (0.0, 0.0, 90.0), atol=1.0e-12) + np.testing.assert_allclose(twist.linear_velocity, (0.0, -math.pi, 0.0), atol=1.0e-12) + omega_rad = np.radians(twist.angular_velocity_deg) + point = np.asarray((3.0, 0.0, 0.0)) + actual_at_point = np.asarray(twist.linear_velocity) + np.cross(omega_rad, point) + expected_at_point = np.cross(omega_rad, point - np.asarray(spec.pivot_point)) + np.testing.assert_allclose(actual_at_point, expected_at_point) + + +def test_curved_twist_requires_an_explicit_radius() -> None: + """A native angular rate cannot be inferred from unspecified task geometry.""" + spec = _surface_spec(curved=True) + + with pytest.raises(ValueError, match="positive radius"): + surface_module.compute_surface_velocity_twist(spec) + + +def test_paths_are_resolved_in_environment_major_order() -> None: + """Runtime rows stay deterministic across stage discovery ordering.""" + specs = (_surface_spec("BeltA"), _surface_spec("Nested/BeltB")) + + paths = surface_module.resolve_surface_velocity_paths(2, specs) + + assert paths == ( + "/World/envs/env_0/BeltA", + "/World/envs/env_0/Nested/BeltB", + "/World/envs/env_1/BeltA", + "/World/envs/env_1/Nested/BeltB", + ) + with pytest.raises(ValueError, match="require every belt"): + surface_module.resolve_surface_velocity_paths(2, (SurfaceVelocitySpec(prim_path="/World/Shared/Belt"),)) + + +def test_facade_ramps_playback_integrates_encoders_and_preserves_commands_on_reset() -> None: + """Full resets restart playback without erasing policy-visible command state.""" + writer = _FakeWriter() + facade = surface_module.SurfaceVelocity(2, (_surface_spec(),), writer=writer) + + assert facade.prim_paths == ("/World/envs/env_0/Belt", "/World/envs/env_1/Belt") + assert facade.num_surfaces == facade.count == 2 + assert [record[2].linear_velocity for record in writer.writes] == [(0.0, 0.0, 0.0)] * 2 + + facade.update(0.25) + np.testing.assert_allclose([record[2].linear_velocity[0] for record in writer.writes[-2:]], (0.1, 0.1)) + np.testing.assert_allclose(facade.get_encoder_positions(), (0.1, 0.1)) + + facade.set_velocities(0.8, indices=[0]) + facade.set_enabled(False, indices=[1]) + facade.update(0.25) + np.testing.assert_allclose(facade.get_commanded_velocities(), (0.8, 0.4)) + np.testing.assert_allclose(facade.get_velocities(), (0.8, 0.0)) + np.testing.assert_allclose(facade.get_encoder_positions(), (0.3, 0.1)) + + facade.reset(env_ids=[0]) + np.testing.assert_allclose(facade.get_encoder_positions(), (0.0, 0.1)) + facade.reset(env_ids=[1, 0]) + np.testing.assert_allclose(facade.get_encoder_positions(), (0.0, 0.0)) + np.testing.assert_allclose(facade.get_commanded_velocities(), (0.8, 0.4)) + assert facade.get_enabled().tolist() == [1, 0] + np.testing.assert_allclose(writer.writes[-2][2].linear_velocity, (0.0, 0.0, 0.0)) + + facade.close() + facade.close() + assert writer.close_count == 1 + assert [(index, enabled) for index, enabled, _ in writer.writes[-2:]] == [(0, False), (1, False)] + + +def test_facade_accepts_torch_control_and_reset_indices() -> None: + """Normal Isaac Lab tensor selectors are copied to host before NumPy validation.""" + facade = surface_module.SurfaceVelocity(2, (_surface_spec(),), writer=_FakeWriter()) + device = "cuda" if torch.cuda.is_available() else "cpu" + + facade.set_velocities(torch.tensor([0.6], device=device), indices=torch.tensor([1], device=device)) + facade.update(0.25) + facade.reset(env_ids=torch.tensor([1], device=device)) + + np.testing.assert_allclose(facade.get_commanded_velocities(), (0.4, 0.6)) + np.testing.assert_allclose(facade.get_encoder_positions(), (0.1, 0.0)) + facade.close() + + +def test_authoring_helper_applies_kinematic_local_surface_schema(monkeypatch: pytest.MonkeyPatch) -> None: + """The lazy authoring seam applies both schemas and authors an initially stopped belt.""" + + class FakeAttribute: + def __init__(self) -> None: + self.value = None + + def Set(self, value) -> bool: + self.value = value + return True + + def Get(self): + return self.value + + class FakePrim: + def __init__(self) -> None: + self.apis = set() + self.attributes = {} + + def IsValid(self) -> bool: + return True + + def HasAPI(self, api_type) -> bool: + return api_type in self.apis + + def GetPath(self) -> str: + return "/World/Belt" + + class FakeRigidBodyAPI: + def __init__(self, prim: FakePrim) -> None: + self.prim = prim + + @classmethod + def Apply(cls, prim: FakePrim): + prim.apis.add(cls) + return cls(prim) + + def CreateRigidBodyEnabledAttr(self) -> FakeAttribute: + return self.prim.attributes.setdefault("rigid_enabled", FakeAttribute()) + + def CreateKinematicEnabledAttr(self) -> FakeAttribute: + return self.prim.attributes.setdefault("kinematic", FakeAttribute()) + + def GetKinematicEnabledAttr(self) -> FakeAttribute: + return self.prim.attributes.setdefault("kinematic", FakeAttribute()) + + class FakeSurfaceAPI: + def __init__(self, prim: FakePrim) -> None: + self.prim = prim + + @classmethod + def Apply(cls, prim: FakePrim): + prim.apis.add(cls) + return cls(prim) + + def _attribute(self, name: str) -> FakeAttribute: + return self.prim.attributes.setdefault(name, FakeAttribute()) + + def CreateSurfaceVelocityLocalSpaceAttr(self) -> FakeAttribute: + return self._attribute("local_space") + + def CreateSurfaceVelocityEnabledAttr(self) -> FakeAttribute: + return self._attribute("surface_enabled") + + def CreateSurfaceVelocityAttr(self) -> FakeAttribute: + return self._attribute("linear") + + def CreateSurfaceAngularVelocityAttr(self) -> FakeAttribute: + return self._attribute("angular") + + fake_pxr = ModuleType("pxr") + fake_pxr.Gf = type("FakeGf", (), {"Vec3f": staticmethod(lambda *values: tuple(values))}) + fake_pxr.PhysxSchema = type("FakePhysxSchema", (), {"PhysxSurfaceVelocityAPI": FakeSurfaceAPI}) + fake_pxr.UsdPhysics = type("FakeUsdPhysics", (), {"RigidBodyAPI": FakeRigidBodyAPI}) + monkeypatch.setitem(sys.modules, "pxr", fake_pxr) + prim = FakePrim() + + surface_module.apply_surface_velocity_api(prim, _surface_spec(), velocity_scale=0.0) + + assert FakeRigidBodyAPI in prim.apis + assert FakeSurfaceAPI in prim.apis + assert prim.attributes["rigid_enabled"].value is True + assert prim.attributes["kinematic"].value is True + assert prim.attributes["local_space"].value is True + assert prim.attributes["surface_enabled"].value is True + assert prim.attributes["linear"].value == (0.0, 0.0, 0.0) + assert prim.attributes["angular"].value == (0.0, 0.0, 0.0) diff --git a/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-franka.minor.rst b/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-franka.minor.rst new file mode 100644 index 000000000000..de2c9432d936 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/maximiliank-conveyor-franka.minor.rst @@ -0,0 +1,21 @@ +Added +^^^^^ + +* Added contributed Franka conveyor tasks with Newton GPU training, CPU-only native PhysX + playback, reusable surface-velocity interfaces, and an interactive Newton-viewer goal selector. +* Added a USD-authored warehouse Play variant with textured 40 mm cartons, gravity infeeds, + compact elevated returns, and 24 physical parcels mapped into the checkpoint's four policy slots. + Added seeded four-color batches, two destination loops, and sorting metrics. Preserved the + original manipulation geometry and 123-observation, eight-action policy interface. Complete-batch + reliability with the unchanged policy was not established. Use ``--viz kit`` for authored visuals + and the explicit base-task checkpoint URL documented in the conveyor guide. +* Added a user guide, preview, and environment-browser entries distinguishing the original + four-cube racetrack task from warehouse sorting with the same pretrained policy. + +Changed +^^^^^^^ + +* Preserved conveyor reset behavior with slice-based environment selections. +* Excluded contacts between stationary warehouse conveyor sections to preserve parcel support after asset replication. +* Kept action-rate penalties finite for rejected NaN or infinite policy commands by tracking the + sanitized commands accepted by the task's action terms. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/README.md b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/README.md new file mode 100644 index 000000000000..82409d287c1e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/README.md @@ -0,0 +1,200 @@ + + +# Conveyor Franka + +Choose between two tasks using the same pretrained Franka policy: + +| Task | Layout and behavior | Newton task ID | +| --- | --- | --- | +| Racetrack transfer | Original two closed racetracks and four numbered cubes; continuous alternating transfers | `IsaacContrib-Conveyor-Franka-Newton-v0` | +| Warehouse sorting | Current extended conveyors and 24 colored parcels; two colors per circulating conveyor | `IsaacContrib-Conveyor-Franka-Newton-Play-v0` | + +The sorter inherits the racetrack environment and configuration. Both reuse the same action, +observation, reward, and placement logic and the same RSL-RL agent configuration. Sorting adds +class dispatch and a four-slot parcel adapter; it needs no separately trained policy. The original +task keeps its compact geometry and four fixed cube identities. Both use 123 observations, +eight actions, 120 Hz physics, and a 60 Hz policy rate. + +## Backend support + +Both tasks run on Newton GPU. `IsaacContrib-Conveyor-Franka-PhysX-CPU-v0` provides a CPU-only +native PhysX reference for the original four-cube racetrack task. + +The PhysX task rejects CUDA during configuration validation. In the supported Isaac Sim runtime, +enabling the native surface-velocity contact-modification path under GPU dynamics can drop the belt +contacts and let packages pass through the conveyor. CPU PhysX preserves those contacts. Use the +Newton task whenever GPU simulation or vectorized throughput is required. + +The two backends deliberately share the policy tensor contract, so an RSL-RL checkpoint can be +loaded by either task without reshaping or reordering tensors. Their contact and actuator dynamics +are not numerically identical; validate task behavior when transferring a policy between them. + +Surface-velocity intent and the tensorized control contract live in +`isaaclab.physics.surface_velocity`. Backend mechanics are separate: Newton's solved-contact force +pipeline lives in `isaaclab_newton.physics.surface_velocity`, while PhysX schema authoring and live +attribute control live in `isaaclab_physx.physics.surface_velocity`. The task package owns only the +racetrack geometry, backend lifecycle selection, and task-level commands. + +## Original racetrack task + +Newton is kitless and supports the lightweight GL viewer: + +```bash +DISPLAY=:1 uv run isaaclab play --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Franka-Newton-v0 \ + --checkpoint pretrained \ + --num_envs 8 --device cuda:0 --viz newton_gl --real-time +``` + +The `pretrained` selector downloads the RSL-RL policy published specifically for the Newton MJWarp +backend. To evaluate another policy, replace `pretrained` with an explicit checkpoint path. The +PhysX task resolves a different backend-specific artifact name, so transferring this Newton policy +to PhysX currently requires the explicit local checkpoint path shown below. + +Training uses the same task ID and defaults to 256 environments: + +```bash +uv run isaaclab train --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Franka-Newton-v0 \ + --num_envs 256 --device cuda:0 +``` + +## Warehouse sorting task + +Select the `Newton-Play` task for warehouse sorting with the same checkpoint. The two parallel manipulation +straights and their adjoining 90-degree bends retain their original positions, widths, radii, +and 0.35 m/s surface speed. Beyond these fixed sections, two short rising feeds climb 0.10 m at less than 10 degrees and +join a shared elevated deck. Guides keep the two return lanes assigned through the upper split, +and separate descending conveyors deliver the parcels back to the original workcell approaches. +The new incline panels use slope-aligned traction and a 0.95 friction coefficient; the original +manipulation sections retain their trained 0.5 setting. + +Blue conveyor frames reach the floor, and the scanner faces along the background main belt after +a 90-degree counterclockwise rotation. The cubes wear SimReady cardboard meshes normalized to +**40 × 40 × 40 mm**, centered on their original colliders; mass remains 50 g. Actions and observation +ordering stay checkpoint-compatible. The Play variant widens its travel bounds and resets all parcels as a randomized mixed batch on the +raised supply belts. +The base training and PhysX tasks retain the original compact layout. + +The workcell contains **24 physical parcels**, exposed to the unchanged checkpoint through +four policy slots. The sorting command fills remote slots with misplaced arrivals; local assignments and an +active grasp stay pinned. Reassignment runs in the command manager, before the next policy observation, +and changes only identity mapping, never a physical pose. Commands, rewards, and placement checks +use the same mapping. All 24 parcels receive belt forces and participate in safety checks. +The arm parks when there is no active sorting transfer. The policy sees remote inventory in canonical waiting slots; +all local parcel poses and velocities, physical transport, rewards, and transfer checks remain +actual simulated states. This adapter is necessary because the checkpoint was trained on compact, +flat returns. Tensor ordering stays 123 observations to eight actions, but remote observation values +are intentionally adapted. Invalid actions still reach the original sanitization and termination path. +The checkpoint can still miss grasps and reset before a batch finishes. The larger randomized +inventory preserves the policy interface, but reliable complete-batch sorting is not established. + +The near loop has a rounded return; the far loop has a shorter squared return with rounded corners +and a shorter supply belt. Both retain the exact original manipulation geometry. Oversized yellow +drive blocks are omitted from the workcell. + +The batch contains six cartons in each of four clearly marked colors: **blue, orange, green, +and purple**. Colored paper bands wrap the textured cardboard without changing its 40 mm bounds. +Reset shuffles physical identities across 24 supply positions, independently +for each environment, using the simulation's seeded random generator. Counts remain balanced while +arrival order and initial conveyor assignments vary. Colors stay fixed throughout each batch. +The 0.043–0.052 m/s feeds release cartons through **12 cm gravity drops** onto the 0.35 m/s loops. +Blue and green belong on the positive-Y loop; orange and purple belong on the negative-Y loop. +The dispatcher requests only wrong-lane transfers and retains ownership through grasp and stable +release. Missed arrivals recirculate for another opportunity. Once the batch is sorted, the arm +parks and both loops continue running. No parcels are recolored, teleported, or replaced during sorting. + +Class assignment is explicit supervisory metadata, not a claim that the unchanged state-based +checkpoint recognizes color. `commands.transfer.parcel_colors` assigns appearances and +`commands.transfer.parcel_destinations` assigns loop IDs; all parcels of a color must share a +single destination. `commands.transfer.randomize_arrivals=False` uses the authored ordering. +`sorted_parcels` counts settled, correctly routed inventory; `batch_complete` indicates that all +24 parcels are settled on their assigned loops. + +The surrounding warehouse includes loaded rack aisles, packing shelves, pallet staging, scan and +outbound signs, safety markings, overhead beams, and warm/cool industrial lighting. Its additional +23.11 m parcel loop includes A29 elevated runs and A38 ramps, with twelve animated cartons traveling +along the main transport line. These background cartons remain render-only inventory and do not +add contacts or policy observations. The 24 small workcell parcels have real collision and +dynamics. A Y-divert dresses the outbound bay. + +Use **Kit/RTX** to see the authored MDL textures, USD lights, and background animation: + +```bash +DISPLAY=:1 uv run --extra isaacsim isaaclab play --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Franka-Newton-Play-v0 \ + --checkpoint https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/6.1/Isaac/IsaacLab/PretrainedCheckpoints/rsl_rl/IsaacContrib-Conveyor-Franka-Newton-v0_newtonmjwarp_none_rsl_rl.pt --num_envs 1 --device cuda:0 --viz kit --real-time \ + --kit_args=--/UJITSO/geometry=false +``` + +Kit renders the USD directly, so the Play configuration excludes visual-only meshes from the +Newton model (`sim.physics.load_visual_shapes=False`). This avoids importing warehouse dressing +into the physics model. For a static approximation in `--viz newton_gl`, explicitly pass +`env.sim.physics.load_visual_shapes=True`; materials and lighting are simplified in that viewer. +The launch override disables experimental geometry streaming to prevent disappearing meshes with +Fabric transforms. Presentation defaults to one environment. Asset references and textures are downloaded and cached +on first use, so the first launch takes longer. + +### Editing the presentation + +* `assets/warehouse.usda` owns the warehouse layout, referenced props, materials, lights, cameras, + and looping parcel transform samples. It can be opened in a USD authoring application with an + Omniverse-compatible asset resolver. `Cameras/Workcell` and `Cameras/Overview` provide two views. +* `assets/parcel_{blue,orange,green,purple}.usda` reference the carton, tint its cardboard, and + add a colored paper band within the original bounds. +* `assets/parcel.usda` normalizes the measured SimReady `cardbox_a1` visual bounds to a centered + 40 mm cube. Keep the shell bounds aligned with the task's original collider when editing it. +* `assets/conveyor_*_supported.usd` override the lower frame vertices of the referenced conveyor + modules. Upper-frame and belt geometry, source topology, and materials remain unchanged. +* `assets/conveyor_routes.usda` owns the elevated network's centerlines, frame meshes, material + references, ramp traction, feed speeds, and mixed-batch spawn positions. Physics reads these same paths after application startup; + USD is not imported during task discovery. Incline normals follow each panel's slope. All 24 + physical parcel spawn positions are authored in `conveyor:parcelSpawnPositions`. +* The Python adapter resolves remote references into the normal local asset cache, removes all + imported physics/action-graph ownership, and samples background animation using simulation time. + `conveyor_warehouse_geometry.py` sweeps lightweight collision proxies along the authored paths. + `mdp/sorting.py` extends the shared transfer command with class dispatch and batch metrics; + the action, observation, reward, and stable-placement contracts remain shared with the trained task. +* `env.conveyor_cube_pool.slot_ids` maps each environment's four policy slots to physical parcel IDs. + `transfer_counts` records placements per physical parcel; assignments never duplicate a parcel + within an environment, and resetting one environment does not change another's slots. + +### Asset and visual references + +The asset survey covered SimReady Central (`simready-central.nvidia.com`), +`omniverse://ov-isaac-dev.nvidia.com/Isaac/SimReady/Industrial/Warehouse`, +`Isaac/Environments/{Digital_Twin_Warehouse,Modular_Warehouse}`, `Isaac/Props/Conveyors`, and +`NVIDIA/Assets/DigitalTwin/Assets/Warehouse`. The composition references publicly accessible +Omniverse counterparts so playback does not require internal Nucleus credentials. Public catalog +usage is described in the [SimReady Explorer documentation](https://docs.omniverse.nvidia.com/extensions/latest/ext_core/ext_browser-extensions/simready-explorer.html). + +Selected assets are the Omniverse A03/A09/A12/A24/A29/A38 conveyors, `RackLarge_A1`, SimReady +`bulkstoragerack_a01` and `cardbox_a1`, the Isaac packing table, and loaded pallets. NVIDIA assets +remain under their original licenses; this package contains the scene layout and lower-frame mesh overrides. +The layout draws on [Dematic's modular conveyor examples](https://www.dematic.com/content/dam/dematic/downloads/brochures/NA_BR_1039_MCS.pdf) +and [KION's warehouse installation photographs](https://www.kiongroup.com/en/News-Stories/Stories/Innovation/Warehousing-is-being-transformed.html): +parallel transport, elevation changes, rack storage, packing zones, and clear marked aisles. +The induction and recirculation layout also follows the concepts in +[Dematic’s sortation overview](https://www.dematic.com/content/dam/dematic/downloads/whitepapers/NA_WP-1015_Sorting-Out-Sortation.pdf). + +## PhysX CPU playback + +The native PhysX variant requires an Isaac Sim-enabled launch and an explicit CPU device. One +environment is the default and recommended interactive configuration: + +```bash +DISPLAY=:1 uv run isaaclab play --rl_library rsl_rl \ + --task IsaacContrib-Conveyor-Franka-PhysX-CPU-v0 \ + --checkpoint /path/to/model.pt \ + --num_envs 1 --device cpu --viz kit --real-time \ + agent.device=cpu +``` + +Overriding the task to CUDA is an error by design; keep the explicit `--device cpu` in launch +commands for clarity. The native surface-velocity backend stages commands through USD, while the +Newton backend keeps its batched state, contact processing, and force application on the GPU. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/__init__.py new file mode 100644 index 000000000000..00b70df752b7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/__init__.py @@ -0,0 +1,44 @@ +# 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 + +"""Four-cube racetrack transfer and warehouse sorting tasks sharing a Franka policy.""" + +import gymnasium as gym + +from . import agents + +gym.register( + id="IsaacContrib-Conveyor-Franka-Newton-v0", + entry_point=f"{__name__}.conveyor_franka_env:ConveyorFrankaEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.conveyor_franka_env_cfg:ConveyorFrankaEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ConveyorFrankaPPORunnerCfg", + }, +) + +gym.register( + # Warehouse sorting reuses the racetrack policy through an explicit checkpoint URL. + id="IsaacContrib-Conveyor-Franka-Newton-Play-v0", + entry_point=f"{__name__}.conveyor_franka_warehouse_env:ConveyorFrankaWarehouseEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.conveyor_franka_asset_env_cfg:ConveyorFrankaA09A12EnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ConveyorFrankaPPORunnerCfg", + }, +) + +gym.register( + # The native PhysxSurfaceVelocityAPI path is intentionally CPU-only. Keep + # that execution contract visible in the public task ID so a CUDA launch is + # never mistaken for a supported configuration. + id="IsaacContrib-Conveyor-Franka-PhysX-CPU-v0", + entry_point=f"{__name__}.conveyor_franka_env:ConveyorFrankaEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.conveyor_franka_physx_env_cfg:ConveyorFrankaPhysxEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ConveyorFrankaPPORunnerCfg", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/__init__.py new file mode 100644 index 000000000000..0e6e80a26022 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Agent configurations for conveyor transfer.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 000000000000..4b11fd71d318 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,182 @@ +# 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 + +"""RSL-RL PPO configuration for conveyor transfer.""" + +import torch +import torch.nn as nn +from rsl_rl.modules.distribution import GaussianDistribution +from torch.distributions import Bernoulli, Normal + +from isaaclab.utils.configclass import configclass + +from isaaclab_rl.rsl_rl import RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg + + +class _ConveyorDeterministicOutput(nn.Module): + """Threshold the gripper logit into the action used by the environment.""" + + def forward(self, output: torch.Tensor) -> torch.Tensor: + gripper = torch.where( + output[..., -1:] >= 0.0, + torch.ones_like(output[..., -1:]), + -torch.ones_like(output[..., -1:]), + ) + return torch.cat((output[..., :-1], gripper), dim=-1) + + +class ConveyorGaussianBernoulliDistribution(GaussianDistribution): + """Use Gaussian exploration for the arm and a Bernoulli gripper. + + The final policy output controls a physically binary action. Sampling it + from a Gaussian would assign different log probabilities to values that + become the same open or close command after thresholding. A Bernoulli + instead optimizes exactly the physical decision seen by the environment. + """ + + def __init__( + self, + output_dim: int, + init_std: float = 0.45, + std_range: tuple[float, float] = (0.15, 0.65), + std_type: str = "scalar", + **kwargs, + ) -> None: + if output_dim < 2: + raise ValueError("The conveyor distribution requires arm outputs followed by one gripper output.") + if len(std_range) != 2 or std_range[0] <= 0.0 or std_range[0] >= std_range[1]: + raise ValueError("std_range must contain positive, increasing bounds.") + if not std_range[0] < init_std < std_range[1]: + raise ValueError("init_std must lie strictly inside std_range.") + if std_type != "scalar": + raise ValueError("The conveyor distribution supports only scalar standard-deviation parameters.") + super().__init__(output_dim, init_std=init_std, std_type=std_type, **kwargs) + self.std_range = (float(std_range[0]), float(std_range[1])) + self._arm_distribution: Normal | None = None + self._gripper_distribution: Bernoulli | None = None + initial_fraction = (init_std - self.std_range[0]) / (self.std_range[1] - self.std_range[0]) + initial_logit = torch.logit(torch.tensor(initial_fraction, dtype=self.std_param.dtype)) + with torch.no_grad(): + self.std_param.fill_(initial_logit) + + def update(self, mlp_output: torch.Tensor) -> None: + """Update continuous-arm and binary-gripper distributions.""" + minimum_std, maximum_std = self.std_range + arm_std = minimum_std + (maximum_std - minimum_std) * torch.sigmoid(self.std_param[:-1]) + self._arm_distribution = Normal(mlp_output[..., :-1], arm_std) + self._gripper_distribution = Bernoulli(logits=mlp_output[..., -1:]) + + def sample(self) -> torch.Tensor: + """Sample seven residuals followed by an exact signed binary action.""" + gripper_open = self._gripper_distribution.sample() + return torch.cat((self._arm_distribution.sample(), 2.0 * gripper_open - 1.0), dim=-1) + + def deterministic_output(self, mlp_output: torch.Tensor) -> torch.Tensor: + """Return arm means and the most likely binary gripper command.""" + return _ConveyorDeterministicOutput()(mlp_output) + + def as_deterministic_output_module(self) -> nn.Module: + """Return an exportable deterministic-output transform.""" + return _ConveyorDeterministicOutput() + + @property + def mean(self) -> torch.Tensor: + """Return arm means and the expected signed gripper command.""" + gripper_mean = 2.0 * self._gripper_distribution.probs - 1.0 + return torch.cat((self._arm_distribution.mean, gripper_mean), dim=-1) + + @property + def std(self) -> torch.Tensor: + """Return arm standard deviations and signed-Bernoulli spread.""" + gripper_std = 2.0 * torch.sqrt(self._gripper_distribution.probs * (1.0 - self._gripper_distribution.probs)) + return torch.cat((self._arm_distribution.stddev, gripper_std), dim=-1) + + @property + def entropy(self) -> torch.Tensor: + """Return joint Gaussian-plus-Bernoulli entropy.""" + return self._arm_distribution.entropy().sum(dim=-1) + self._gripper_distribution.entropy().sum(dim=-1) + + @property + def params(self) -> tuple[torch.Tensor, ...]: + """Return parameters needed to evaluate the mixed KL divergence.""" + return self._arm_distribution.mean, self._arm_distribution.stddev, self._gripper_distribution.logits + + def log_prob(self, outputs: torch.Tensor) -> torch.Tensor: + """Evaluate the exact continuous/binary physical action.""" + arm_log_prob = self._arm_distribution.log_prob(outputs[..., :-1]).sum(dim=-1) + gripper_open = (outputs[..., -1:] >= 0.0).to(outputs.dtype) + return arm_log_prob + self._gripper_distribution.log_prob(gripper_open).sum(dim=-1) + + def kl_divergence( + self, + old_params: tuple[torch.Tensor, ...], + new_params: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + """Return ``KL(old || new)`` for both action families.""" + old_arm_mean, old_arm_std, old_gripper_logits = old_params + new_arm_mean, new_arm_std, new_gripper_logits = new_params + arm_kl = torch.distributions.kl_divergence( + Normal(old_arm_mean, old_arm_std), + Normal(new_arm_mean, new_arm_std), + ).sum(dim=-1) + old_gripper_probability = torch.sigmoid(old_gripper_logits) + gripper_kl = ( + old_gripper_probability * (old_gripper_logits - new_gripper_logits) + - torch.nn.functional.softplus(old_gripper_logits) + + torch.nn.functional.softplus(new_gripper_logits) + ).sum(dim=-1) + return arm_kl + gripper_kl.clamp_min(0.0) + + +@configclass +class ConveyorGaussianBernoulliDistributionCfg(RslRlMLPModelCfg.GaussianDistributionCfg): + """Bounded Gaussian arm exploration with a Bernoulli gripper.""" + + class_name: str = ( + "isaaclab_tasks.contrib.conveyor_franka.agents.rsl_rl_ppo_cfg:ConveyorGaussianBernoulliDistribution" + ) + std_range: tuple[float, float] = (0.15, 0.65) + + +@configclass +class ConveyorFrankaPPORunnerCfg(RslRlOnPolicyRunnerCfg): + """PPO configuration for four-cube commanded transfer.""" + + num_steps_per_env = 32 + # RSL-RL's usual randomized episode counters would desynchronize the + # per-subgoal timeout clock before the first policy step. + init_at_random_ep_len = False + max_iterations = 4000 + save_interval = 50 + experiment_name = "conveyor_franka_transfer" + clip_actions = 1.0 + obs_groups = {"actor": ["policy"], "critic": ["policy"]} + actor = RslRlMLPModelCfg( + hidden_dims=[512, 256, 128], + activation="elu", + obs_normalization=True, + distribution_cfg=ConveyorGaussianBernoulliDistributionCfg(init_std=0.45), + ) + critic = RslRlMLPModelCfg( + hidden_dims=[512, 256, 128], + activation="elu", + obs_normalization=True, + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.001, + num_learning_epochs=5, + num_mini_batches=16, + learning_rate=1.0e-4, + schedule="fixed", + # At 60 Hz, the ten-second pickup-to-placement horizon needs the same + # long-horizon discount used by the reset-driven Franka stack task. + gamma=0.999, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_quarter_supported.usd b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_quarter_supported.usd new file mode 100644 index 000000000000..2e3ce9e849d4 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_quarter_supported.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:281020b9d284f23e224babda0210629c64f663a74482a7962d1c511ba73bce8a +size 7997402 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_routes.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_routes.usda new file mode 100644 index 000000000000..c35390b0218f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_routes.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ffa5fa3f772265e0d6e563b5ffdf3ea315602dd45ce4a2205a30477aee6e572 +size 1045928 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_straight_supported.usd b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_straight_supported.usd new file mode 100644 index 000000000000..b07af0f62a85 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/conveyor_straight_supported.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4a735b6c8cb283e4cda2950d2e4ef6db6b2e91048e3d4b4a4d5d285fa806e1a +size 22319911 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel.usda new file mode 100644 index 000000000000..994effb5c4bd --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9f8366dd61f93ade13f11d0e193ecd389c94f4e9493de0d97f9a2ff908b6156d +size 890 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_blue.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_blue.usda new file mode 100644 index 000000000000..ff92e414ee8b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_blue.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:53894664629ad77181aac1416482a613a782e261b0dd801e36d2a7e5fedd01e2 +size 1926 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_green.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_green.usda new file mode 100644 index 000000000000..eb7320c43338 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_green.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:371036df39c3586de1166ad93753fab1bc45d9d8de8bd0b49077a29ad16287aa +size 1928 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_orange.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_orange.usda new file mode 100644 index 000000000000..de6de478fdf1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_orange.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5076a92ba6b4141f7a9c64ecb7b2789fb2058f862f933c2b2217fc4696c8a9b8 +size 1927 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_purple.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_purple.usda new file mode 100644 index 000000000000..750fe3441703 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/parcel_purple.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8a39e3d6f918d99cf8a811c68161565c224da9ad8ab2889ebd08e7232807bd38 +size 1924 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/warehouse.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/warehouse.usda new file mode 100644 index 000000000000..c5a87bb58e21 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/assets/warehouse.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1b2f7d68aa2aaddf17a67c7f44fe728456b4ce23ae6d388a6fa557c8657dcdb3 +size 458277 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_cube_pool.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_cube_pool.py new file mode 100644 index 000000000000..2f241455277c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_cube_pool.py @@ -0,0 +1,99 @@ +# 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 + +"""Stable physical parcel identities behind the checkpoint's four observation slots.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from .mdp.reset_events import CUBE_COUNT + +if TYPE_CHECKING: + from isaaclab.assets import RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +def cube_values(env: ManagerBasedRLEnv, attribute: str, *, all_cubes: bool = False) -> torch.Tensor: + """Gather a rigid-object data attribute in policy-slot or complete physical-inventory order. + + Values retain the attribute's units and world frame. The result has shape + ``(num_envs, num_slots_or_parcels, attribute_size)``. Tasks without a parcel pool + retain the original four fixed cube identities. + """ + pool = getattr(env, "conveyor_cube_pool", None) + assets = pool.assets if pool is not None else tuple(env.scene[f"cube_{i}"] for i in range(CUBE_COUNT)) + values = torch.stack([getattr(asset.data, attribute).torch for asset in assets], dim=1) + if pool is not None and not all_cubes: + values = values.gather(1, pool.slot_ids[..., None].expand(-1, -1, values.shape[-1])) + return values + + +class ConveyorCubePool: + """Bind four policy slots to a larger pool without moving or copying physical bodies. + + Assignments are independent per environment. Local parcels keep their slots; + the active grasp is additionally pinned even if it moves outside the workcell. + """ + + def __init__(self, assets: tuple[RigidObject, ...], num_envs: int, device: str) -> None: + if len(assets) < CUBE_COUNT: + raise ValueError("The conveyor policy requires at least four physical parcels.") + self.assets = assets + self.slot_ids = torch.arange(CUBE_COUNT, device=device).repeat(num_envs, 1) + self.assignment_counts = torch.zeros((num_envs, len(assets)), device=device, dtype=torch.long) + self.assignment_counts[:, :CUBE_COUNT] = 1 + self.transfer_counts = torch.zeros_like(self.assignment_counts) + + def reset(self, env_ids: Sequence[int] | torch.Tensor) -> None: + """Restore reset-recipe slot identities for selected environments.""" + self.slot_ids[env_ids] = torch.arange(CUBE_COUNT, device=self.slot_ids.device) + + def refresh( + self, + positions: torch.Tensor, + local: torch.Tensor, + candidates: torch.Tensor, + target_slots: torch.Tensor, + pinned: torch.Tensor, + ) -> torch.Tensor: + """Assign arriving parcels to remote slots, preserving all local and pinned identities. + + Args: + positions: Physical parcel positions in the workspace [m], shape ``(N, P, 3)``. + local: Parcels still in the manipulation region, shape ``(N, P)``. + candidates: Parcels eligible for pickup, shape ``(N, P)``. + target_slots: Current command's policy slot, shape ``(N,)``. + pinned: Whether the active parcel must retain its slot, shape ``(N,)``. + + Returns: + Environments whose slot assignments changed, shape ``(N,)``. + """ + mapped = torch.zeros_like(candidates).scatter_(1, self.slot_ids, True) + available = candidates & ~mapped + reusable = ~local.gather(1, self.slot_ids) + reusable.scatter_(1, target_slots[:, None], ~pinned[:, None] & reusable.gather(1, target_slots[:, None])) + changed = torch.zeros_like(pinned) + # Prefer parcels assigned less often, then the arrival with more belt travel remaining. + priority = positions[..., 0] - 10.0 * self.assignment_counts + for slot in range(CUBE_COUNT): + eligible = reusable[:, slot] & available.any(dim=1) + rows = eligible.nonzero(as_tuple=False).flatten() + if not rows.numel(): + continue + choices = torch.where(available, priority, -torch.inf).argmax(dim=1)[rows] + self.slot_ids[rows, slot] = choices + self.assignment_counts[rows, choices] += 1 + available[rows, choices] = False + changed[rows] = True + return changed + + def record_transfers(self, env_ids: torch.Tensor, target_slots: torch.Tensor) -> None: + """Credit stable placements to physical parcels before the command selects its next slot.""" + physical_ids = self.slot_ids[env_ids, target_slots] + self.transfer_counts[env_ids, physical_ids] += 1 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_env_cfg.py new file mode 100644 index 000000000000..899f068f5090 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_env_cfg.py @@ -0,0 +1,399 @@ +# 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 + +"""Digital Twin warehouse visuals for conveyor-Franka playback.""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import cache +from pathlib import Path +from typing import TYPE_CHECKING + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.physics import SurfaceVelocitySpec +from isaaclab.terrains import TerrainImporterCfg +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, retrieve_file_path +from isaaclab.utils.configclass import configclass + +if TYPE_CHECKING: + from pxr import Sdf, Usd + + from isaaclab.terrains import TerrainImporter + +from .conveyor_franka_env_cfg import ( + ConveyorFrankaEnvCfg, + ConveyorFrankaSceneCfg, + _hidden_collision_geometry, + _hidden_collision_mesh, + _spawn_shape_with_display_color, +) +from .conveyor_geometry import ( + BELT_CENTER_X, + BELT_CENTER_Y, + BELT_HALF_STRAIGHT, + BELT_TOP_Z, + BELT_TURN_RADIUS, +) +from .conveyor_warehouse_geometry import warehouse_belt_sections, warehouse_guard_meshes + +_THOR_TABLE_ASSET_PATH = f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/thor_table.usd" + +# The A12 endpoints are 2.9922 m apart, its belt crown is 1.78053 m above the +# asset origin. USD point overrides extend the legs to the floor while keeping +# the upper frames unchanged; the policy workspace follows the deck elevation. +_A12_ENDPOINT_SEPARATION = 2.9922 +_ASSET_BELT_TOP_Z = 1.78053 +_ASSET_XY_SCALE = 2.0 * BELT_TURN_RADIUS / _A12_ENDPOINT_SEPARATION +_CONVEYOR_SUPPORT_Z = 0.55 +# Preserve the upper frames' proportions. The scaled belt crown determines +# how far to elevate the policy workspace. +_ASSET_Z_SCALE = _ASSET_XY_SCALE +_ASSET_ROOT_Z = _CONVEYOR_SUPPORT_Z +_ASSET_BELT_WORLD_Z = _ASSET_ROOT_Z + _ASSET_BELT_TOP_Z * _ASSET_Z_SCALE +_WORKSPACE_ELEVATION = _ASSET_BELT_WORLD_Z - BELT_TOP_Z + +# A09 is a 4 m straight. Its travel-axis scale is independent of the common +# lateral scale so one asset spans the task's complete 0.88 m straight run. +_A09_LENGTH = 4.0 +_A09_X_SCALE = 2.0 * BELT_HALF_STRAIGHT / _A09_LENGTH + +# The Thor table authors its mounting surface at local z=0 and its lowest foot +# at z=-0.795 m. Uniformly scaling that distance to the elevated robot base +# puts every foot on the global ground without distorting the table. +_THOR_TABLE_LOWEST_Z = -0.795 +_THOR_TABLE_SCALE = _WORKSPACE_ELEVATION / -_THOR_TABLE_LOWEST_Z + + +_PHYSICS_SCHEMA_PREFIXES = ("Physics", "Physx", "Newton", "Mujoco") +_PHYSICS_SCHEMA_NAMES = frozenset(("IsaacConveyorAPI",)) +_PRESENTATION_ASSETS = Path(__file__).parent / "assets" + + +@cache +def _presentation_layer(usd_path: str) -> Sdf.Layer: + """Resolve a local composition's remote assets without modifying its source layer.""" + from pxr import Sdf, UsdUtils + + source = Sdf.Layer.FindOrOpen(usd_path) + layer = Sdf.Layer.CreateAnonymous(Path(usd_path).name) + layer.TransferContent(source) + resolved = {} + for path in source.GetExternalReferences(): + if path.startswith("https://"): + resolved[path] = retrieve_file_path(path) + else: + local_path = Sdf.ComputeAssetPathRelativeToLayer(source, path) + resolved[path] = _presentation_layer(local_path).identifier + UsdUtils.ModifyAssetPaths( + layer, + lambda path: resolved[path] if path in resolved else Sdf.ComputeAssetPathRelativeToLayer(source, path), + ) + return layer + + +@sim_utils.clone +def _spawn_authored_visual( + prim_path: str, + cfg: sim_utils.UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Compose USD scenery with cached dependencies and no simulation ownership.""" + layer = _presentation_layer(cfg.usd_path) + prim = sim_utils.create_prim( + prim_path, + translation=translation, + orientation=orientation, + scale=cfg.scale, + ) + prim.GetReferences().AddReference(layer.identifier) + sim_utils.make_uninstanceable(prim_path) + _make_usd_subtree_visual_only(prim) + return prim + + +@sim_utils.clone +def _spawn_carton_cube( + prim_path: str, + cfg: _ParcelCuboidCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Dress the original 40 mm collider with a centered, equally sized SimReady carton.""" + from pxr import UsdGeom + + prim = _spawn_shape_with_display_color(prim_path, cfg, translation, orientation, **kwargs) + UsdGeom.Imageable(prim.GetStage().GetPrimAtPath(f"{prim_path}/geometry/mesh")).MakeInvisible() + visual_cfg = sim_utils.UsdFileCfg(usd_path=cfg.parcel_usd_path) + _spawn_authored_visual(f"{prim_path}/CartonVisual", visual_cfg) + return prim + + +@configclass +class _ParcelCuboidCfg(sim_utils.CuboidCfg): + """Original task collider with a separately authored carton appearance.""" + + parcel_usd_path: str = str(_PRESENTATION_ASSETS / "parcel.usda") + """USD visual, normalized to the task's 40 mm cube.""" + + +def _is_physics_schema(schema_name: str) -> bool: + """Return whether an applied schema can add physics ownership to a visual asset.""" + return schema_name in _PHYSICS_SCHEMA_NAMES or schema_name.startswith(_PHYSICS_SCHEMA_PREFIXES) + + +def _make_usd_subtree_visual_only(root_prim) -> None: + """Author a render-only override for a referenced USD subtree. + + The source asset remains untouched. All edits are stronger opinions in the + task stage and fail closed if a physics schema cannot be removed. + """ + # Import USD only when a stage exists. Besides respecting Kit startup + # ordering, this keeps import-light task discovery kitless. + from pxr import Sdf, Usd, UsdPhysics + + children = tuple(Usd.PrimRange(root_prim, Usd.TraverseInstanceProxies())) + instance_proxies = tuple(str(child.GetPath()) for child in children if child.IsInstanceProxy()) + if instance_proxies: + raise RuntimeError( + "Visual-only USD overrides require editable descendants; set make_uninstanceable=True. " + f"Found instance proxies below {root_prim.GetPath()}: {instance_proxies[:3]}" + ) + + with Sdf.ChangeBlock(): + for child in children: + if child.GetTypeName().startswith("OmniGraph") or child.IsA(UsdPhysics.Scene): + child.SetActive(False) + + for child in children: + if not child.IsValid() or not child.IsActive(): + continue + if child.IsA(UsdPhysics.Joint): + child.SetActive(False) + continue + for schema_name in tuple(child.GetAppliedSchemas()): + if _is_physics_schema(schema_name) and not child.RemoveAppliedSchema(schema_name): + raise RuntimeError(f"Failed to remove physics schema {schema_name!r} from {child.GetPath()}.") + + remaining = { + str(child.GetPath()): tuple(schema for schema in child.GetAppliedSchemas() if _is_physics_schema(schema)) + for child in children + if child.IsValid() and child.IsActive() + } + remaining = {path: schemas for path, schemas in remaining.items() if schemas} + if remaining: + raise RuntimeError(f"Visual-only USD subtree still contains physics schemas: {remaining}") + + +@sim_utils.clone +def _spawn_visual_only_usd( + prim_path: str, + cfg: sim_utils.UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +): + """Spawn a USD asset and strip its authored physics metadata.""" + prim = sim_utils.spawn_from_usd(prim_path, cfg, translation, orientation, **kwargs) + + # Presentation assets may contain nested dynamic props (the packing table, + # for example, carries an authored container rigid body). Merely disabling + # collision still exposes those bodies and joints to a backend parser. + # Author local API deletions so the entire referenced hierarchy is a pure + # render layer and cannot alter either Newton's model or PhysX's scene. + _make_usd_subtree_visual_only(prim) + return prim + + +@configclass +class _VisualOnlyUsdFileCfg(sim_utils.UsdFileCfg): + """USD reference whose composed subtree is guaranteed to remain render-only.""" + + func: Callable = _spawn_visual_only_usd + # Recursive schema overrides cannot be authored on USD instance proxies. + make_uninstanceable: bool = True + + +@configclass +class _ElevatedGroundPlaneCfg(TerrainImporterCfg): + """Ground plane that reports the elevated workspace as its environment origin.""" + + class_type: type[TerrainImporter] | str = ( + "{DIR}.conveyor_franka_asset_terrain:ConveyorFrankaGroundPlaneTerrainImporter" + ) + workspace_origin_offset: tuple[float, float, float] = (0.0, 0.0, _WORKSPACE_ELEVATION) + """Translation from clone-grid origins to policy workspaces [m].""" + + +def _visual_usd_asset( + prim_path: str, + usd_path: str, + position: tuple[float, float, float], + scale: tuple[float, float, float], + rotation: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), +) -> AssetBaseCfg: + """Build one visual asset whose authored colliders are explicitly disabled.""" + spawn = _VisualOnlyUsdFileCfg( + usd_path=usd_path, + scale=scale, + ) + return AssetBaseCfg( + prim_path=prim_path, + init_state=AssetBaseCfg.InitialStateCfg(pos=position, rot=rotation), + # Physics remains owned by the task's lightweight, hidden belt and rail + # proxies. The spawn callback strips authored physics metadata to keep + # the visual geometry entirely non-authoritative. + spawn=spawn, + ) + + +@configclass +class ConveyorFrankaA09A12SceneCfg(ConveyorFrankaSceneCfg): + """Checkpoint-compatible Digital Twin scene with visual warehouse dressing.""" + + def __post_init__(self) -> None: + """Dress the workcell and extend its returns around the fixed manipulation sections.""" + super().__post_init__() + + # Raise every inherited env-scoped task component as one rigid + # workspace. The environment reports the same offset as its origin, so + # observations, resets, rewards, and the pretrained policy retain their + # original local coordinates. + for asset in vars(self).values(): + if isinstance(asset, AssetBaseCfg) and asset.prim_path.startswith("{ENV_REGEX_NS}/"): + x, y, z = asset.init_state.pos + asset.init_state.pos = (x, y, z + _WORKSPACE_ELEVATION) + + # The ground is global rather than env-scoped and remains at the USD + # convention's default elevation. + self.ground = _ElevatedGroundPlaneCfg( + prim_path="/World/GroundPlane", + terrain_type="plane", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.055, 0.065, 0.082), roughness=0.82), + ) + self.dome_light.spawn.color = (0.70, 0.78, 0.92) + self.dome_light.spawn.intensity = 700.0 + + right_x = BELT_CENTER_X + BELT_HALF_STRAIGHT + straight_scale = (_A09_X_SCALE, _ASSET_XY_SCALE, _ASSET_Z_SCALE) + + # The Franka is fixed at the elevated workspace origin and does not need + # support collision. Replace the temporary plinth with the purpose-built + # Thor table: its mount stays at the robot base and its feet land on z=0. + self.tabletop = _visual_usd_asset( + prim_path="{ENV_REGEX_NS}/RobotThorTableVisual", + usd_path=_THOR_TABLE_ASSET_PATH, + position=(0.0, 0.0, _WORKSPACE_ELEVATION), + scale=(_THOR_TABLE_SCALE,) * 3, + ) + self.table_pedestal = None + + for side in ("Left", "Right"): + side_key = side.lower() + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + + # The Digital Twin pieces already render their belt, frame, and + # guides, so remove only the procedural render geometry. Hidden + # belt and guide collision assets created above stay authoritative. + setattr(self, f"conveyor_{side_key}_belt_visual", None) + setattr(self, f"guard_{side_key}_inner_visual", None) + setattr(self, f"guard_{side_key}_outer_visual", None) + + run = "bottom" if side == "Left" else "top" + y_position = center_y + (-BELT_TURN_RADIUS if side == "Left" else BELT_TURN_RADIUS) + setattr( + self, + f"conveyor_{side_key}_{run}_a09_visual", + _visual_usd_asset( + prim_path=f"{{ENV_REGEX_NS}}/Conveyor{side}{run.title()}A09Visual", + usd_path=str(_PRESENTATION_ASSETS / "conveyor_straight_supported.usd"), + position=(right_x, y_position, _ASSET_ROOT_Z), + scale=straight_scale, + ), + ) + getattr(self, f"conveyor_{side_key}_{run}_a09_visual").spawn.func = _spawn_authored_visual + + self.warehouse_visual = _visual_usd_asset( + prim_path="{ENV_REGEX_NS}/WarehouseVisual", + usd_path=str(_PRESENTATION_ASSETS / "warehouse.usda"), + position=(0.0, 0.0, 0.0), + scale=(1.0, 1.0, 1.0), + ) + self.warehouse_visual.spawn.func = _spawn_authored_visual + for cube_id in range(4): + cube = getattr(self, f"cube_{cube_id}") + cube.spawn = _ParcelCuboidCfg(**vars(cube.spawn)) + cube.spawn.func = _spawn_carton_cube + + def _configure_route_assets( + self, parcel_colors: tuple[str, ...] = ("blue", "orange", "green", "purple") * 6 + ) -> None: + """Read USD route geometry after the application selects its USD runtime.""" + from .conveyor_warehouse_geometry import warehouse_parcel_positions + + positions = warehouse_parcel_positions() + if len(parcel_colors) != len(positions) or not set(parcel_colors) <= {"blue", "orange", "green", "purple"}: + raise ValueError("Each authored parcel requires a color: blue, orange, green, or purple.") + for cube_id, (position, color) in enumerate(zip(positions, parcel_colors, strict=True)): + cube = self.cube_0.copy() + cube.spawn.parcel_usd_path = str(_PRESENTATION_ASSETS / f"parcel_{color}.usda") + cube.prim_path = f"{{ENV_REGEX_NS}}/Cube{cube_id}" + cube.init_state.pos = (position[0], position[1], position[2] + _WORKSPACE_ELEVATION) + setattr(self, f"cube_{cube_id}", cube) + for side in ("Left", "Right"): + for key in ("top_straight", "bottom_straight", "right_turn", "left_turn"): + setattr(self, f"conveyor_{side.lower()}_{key}_collision", None) + for index, section in enumerate(warehouse_belt_sections(side)): + asset = _hidden_collision_geometry(section.belt.prim_path, section.geometry, 1.1e-5, 1) + x, y, z = asset.init_state.pos + asset.init_state.pos = (x, y, z + _WORKSPACE_ELEVATION) + setattr(self, f"warehouse_{side.lower()}_section_{index}", asset) + for boundary in ("inner", "outer"): + setattr(self, f"guard_{side.lower()}_{boundary}_collision", None) + for guard in warehouse_guard_meshes(side): + asset = _hidden_collision_mesh(f"{{ENV_REGEX_NS}}/{guard.name}Collision", guard, 1.1e-5, 1) + asset.init_state.pos = (0.0, 0.0, _WORKSPACE_ELEVATION) + setattr(self, f"warehouse_{guard.name}_collision", asset) + + def build_conveyor_belt_specs(self, **kwargs: float | bool) -> tuple[SurfaceVelocitySpec, ...]: + """Describe the linked warehouse surfaces to the existing Newton conveyor driver.""" + return tuple(section.belt for side in ("Left", "Right") for section in warehouse_belt_sections(side, **kwargs)) + + +@configclass +class ConveyorFrankaA09A12EnvCfg(ConveyorFrankaEnvCfg): + """Newton presentation variant with Digital Twin visuals and extended return routes.""" + + scene: ConveyorFrankaA09A12SceneCfg = ConveyorFrankaA09A12SceneCfg( + num_envs=1, + env_spacing=24.0, + replicate_physics=True, + ) + + def __post_init__(self) -> None: + """Frame the presentation scene and allow packages to travel along its extended returns.""" + super().__post_init__() + from .mdp.sorting import ConveyorSortCommandCfg + + self.commands.transfer = ConveyorSortCommandCfg() + # Kit renders the authored USD directly; Newton needs only the physical geometry. + self.sim.physics.load_visual_shapes = False + # Adjacent static belt sections must not consume the parcel contact budget. + self.sim.physics.collision_cfg.include_static_kinematic_pairs = False + self.conveyor_force.transported_body_count_per_env = len(self.commands.transfer.parcel_destinations) + self.sim.physics.solver_cfg.nconmax = 400 + self.sim.physics.solver_cfg.njmax = 600 + self.conveyor_force.transported_body_pattern = r"(?:^|/)Cube_?[0-9]+(?:/|$)" + self.terminations.cube_out_of_workspace.params = { + "minimum": (-0.4, -1.2, -0.05), + "maximum": (2.85, 1.2, 0.8), + } + self.sim.default_visualizer_cfg.eye = (4.8, -5.2, 3.0) + self.sim.default_visualizer_cfg.lookat = (0.9, 0.6, 0.95) + self.sim.default_visualizer_cfg.focal_length = 24.0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_terrain.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_terrain.py new file mode 100644 index 000000000000..645b58511d27 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_asset_terrain.py @@ -0,0 +1,30 @@ +# 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 + +"""Terrain-origin support for the asset-rich conveyor playback variant.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.terrains import TerrainImporter + +if TYPE_CHECKING: + from .conveyor_franka_asset_env_cfg import _ElevatedGroundPlaneCfg + + +class ConveyorFrankaGroundPlaneTerrainImporter(TerrainImporter): + """Plane terrain whose environment origins follow an elevated workspace.""" + + def __init__(self, cfg: _ElevatedGroundPlaneCfg): + """Create the ground plane and translate its policy-facing origins.""" + super().__init__(cfg) + self.env_origins.add_(self.env_origins.new_tensor(cfg.workspace_origin_offset)) + # The authored warehouse floor replaces the default calibration-grid visual. + from pxr import UsdGeom + + from isaaclab.sim.utils.stage import get_current_stage + + UsdGeom.Imageable(get_current_stage().GetPrimAtPath(cfg.prim_path)).MakeInvisible() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env.py new file mode 100644 index 000000000000..7e75d4835b21 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env.py @@ -0,0 +1,134 @@ +# 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 + +"""Manager-based environment that installs the selected conveyor physics adapter.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.physics import SurfaceVelocityView + +from .conveyor_franka_env_cfg import ConveyorFrankaEnvCfg +from .conveyor_geometry import belt_collision_section_specs +from .conveyor_goal_selector import ConveyorGoalSelector + + +class ConveyorFrankaEnv(ManagerBasedRLEnv): + """Manager-based environment with backend-native conveyor surfaces.""" + + cfg: ConveyorFrankaEnvCfg + + def __init__(self, cfg: ConveyorFrankaEnvCfg, render_mode: str | None = None, **kwargs): + self._conveyor_driver: SurfaceVelocityView | None = None + super().__init__(cfg, render_mode=render_mode, **kwargs) + self._goal_selector: ConveyorGoalSelector | None = None + self._setup_goal_selector() + + def _init_sim(self) -> None: + """Install the conveyor adapter at the lifecycle point required by its backend.""" + belt_spec_kwargs = { + "velocity": self.cfg.conveyor_force.speed, + "friction_coefficient": self.cfg.conveyor_force.friction, + "contact_threshold": self.cfg.conveyor_force.normal_threshold, + } + spec_builder = getattr(self.cfg.scene, "build_conveyor_belt_specs", None) + if spec_builder is None: + belt_specs = tuple( + section.belt + for side in ("Left", "Right") + for section in belt_collision_section_specs(side, **belt_spec_kwargs) + ) + else: + belt_specs = tuple(spec_builder(**belt_spec_kwargs)) + env_path_format = self.cfg.scene.clone_cfg.clone_template + + # Newton needs solved-contact attributes and graph callbacks registered + # before the first reset finalizes and captures the solver. PhysX belt + # schemas, by contrast, are authored by the scene spawners and its live + # command adapter is attached only after PhysX has parsed that scene. + from isaaclab_newton.physics import NewtonCfg, SurfaceVelocity + + if isinstance(self.cfg.sim.physics, NewtonCfg): + driver = SurfaceVelocity( + num_envs=self.cfg.scene.num_envs, + surface_specs=belt_specs, + startup_duration_s=self.cfg.conveyor_force.startup_duration_s, + env_path_format=env_path_format, + body_pattern=self.cfg.conveyor_force.transported_body_pattern, + body_count_per_env=self.cfg.conveyor_force.transported_body_count_per_env, + ) + self._conveyor_driver = driver + try: + super()._init_sim() + except Exception: + driver.close() + self._conveyor_driver = None + raise + return + + from isaaclab_physx.physics import PhysxCfg, SurfaceVelocity + + if not isinstance(self.cfg.sim.physics, PhysxCfg): + raise ValueError(f"Unsupported conveyor physics backend: {type(self.cfg.sim.physics).__name__}.") + + configure_conveyor = getattr(self.cfg.scene, "configure_conveyor", None) + if configure_conveyor is not None: + configure_conveyor(friction_coefficient=self.cfg.conveyor_force.friction) + super()._init_sim() + driver = SurfaceVelocity( + num_envs=self.cfg.scene.num_envs, + surface_specs=belt_specs, + env_path_format=env_path_format, + startup_duration_s=self.cfg.conveyor_force.startup_duration_s, + stage=self.sim.stage, + ) + try: + driver.start() + except Exception: + driver.close() + raise + self._conveyor_driver = driver + + @property + def conveyor_belt(self) -> SurfaceVelocityView: + """Tensorized conveyor control view for this environment.""" + driver = self._conveyor_driver + if driver is None: + raise RuntimeError("The conveyor belt is unavailable before simulation initialization or after close().") + return driver + + def _setup_goal_selector(self) -> None: + """Attach one task panel to the first supported interactive visualizer.""" + for visualizer in self.sim.visualizers: + if getattr(visualizer.cfg, "visualizer_type", None) not in {"newton_gl", "newton_rtx"}: + continue + register_callback = getattr(visualizer, "register_ui_callback", None) + if register_callback is None: + continue + visible_env_ids = visualizer.get_visualized_env_ids() + env_id = visible_env_ids[0] if visible_env_ids else 0 + self._goal_selector = ConveyorGoalSelector(self, env_id) + register_callback(self._goal_selector.render, position="panel") + return + + def _reset_idx(self, env_ids: Sequence[int] | slice): + """Reset selected environments and discard stale conveyor forces.""" + if isinstance(env_ids, slice): + env_ids = self.scene._ALL_INDICES[env_ids] + super()._reset_idx(env_ids) + + conveyor_driver = getattr(self, "_conveyor_driver", None) + if conveyor_driver is not None: + conveyor_driver.reset(env_ids) + + def close(self): + """Release conveyor callbacks before the physics scene is destroyed.""" + conveyor_driver = getattr(self, "_conveyor_driver", None) + if conveyor_driver is not None: + conveyor_driver.close() + self._conveyor_driver = None + super().close() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env_cfg.py new file mode 100644 index 000000000000..e6e1acf19d76 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_env_cfg.py @@ -0,0 +1,725 @@ +# 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 + +"""Configuration for the force-driven conveyor and Franka demonstration scene.""" + +from __future__ import annotations + +import math +import re + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg, NewtonCollisionPipelineCfg, NewtonShapeCfg +from isaaclab_newton.sim.schemas import MujocoCollisionCfg, NewtonCollisionCfg, NewtonMaterialPropertiesCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import CurriculumTermCfg as CurrTerm +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.sim.schemas import CollisionFragment, UsdPhysicsCollisionCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.utils.configclass import configclass +from isaaclab.visualizers import VisualizerCfg + +from . import mdp +from .conveyor_geometry import ( + BELT_COLOR, + BELT_INNER_STRAIGHT_Y, + BELT_OUTER_STRAIGHT_Y, + CUBE_COLORS, + CUBE_INNER_SLOT_X, + CUBE_OUTER_SLOT_X, + GUARD_COLOR, + CuboidSpec, + MeshSpec, + belt_collision_geometry_specs, + belt_mesh_spec, + guard_mesh_specs, +) +from .franka_robot_cfg import FRANKA_PANDA_CONVEYOR_CFG +from .mdp.terminations import invalid_action as invalid_policy_action + +_DYNAMIC_PROPERTIES = sim_utils.RigidBodyBaseCfg() +_CONTACT_GAP = 0.01 +_CUBE_CONTACT_MARGIN = 0.003 +_MUJOCO_SOLIMP = (0.9, 0.95, 0.001, 0.5, 2.0) +_MUJOCO_SOLREF = (0.02, 1.0) +_SUBGOAL_TIMEOUT_S = 20.0 +_TRANSFER_SEQUENCE_LENGTH = 8 +_ARM_JOINT_NAMES = tuple(f"panda_joint{joint_id}" for joint_id in range(1, 8)) +_FINGER_JOINT_NAMES = ("panda_finger_joint1", "panda_finger_joint2") + + +def _validate_common_config(cfg: ConveyorFrankaEnvCfg) -> None: + """Validate backend-independent timing and policy tensor contracts.""" + cfg.conveyor_force.validate_config() + if not cfg.sim.use_newton_actuators: + raise ValueError("The conveyor Franka requires the shared Newton-actuator execution path.") + if not math.isfinite(cfg.sim.dt) or cfg.sim.dt <= 0.0 or cfg.decimation <= 0: + raise ValueError("Simulation dt and environment decimation must be positive.") + arm_action = cfg.actions.arm_action + if arm_action.joint_names != list(_ARM_JOINT_NAMES) or not arm_action.preserve_order: + raise ValueError("Arm actions must preserve the explicit panda_joint1-to-panda_joint7 ordering.") + if not math.isfinite(arm_action.max_delta) or arm_action.max_delta <= 0.0: + raise ValueError("Arm max_delta must be finite and positive.") + if not math.isfinite(arm_action.joint_limit_margin) or arm_action.joint_limit_margin < 0.0: + raise ValueError("Arm joint_limit_margin must be finite and non-negative.") + lower = arm_action.workspace_lower + upper = arm_action.workspace_upper + if len(lower) != len(_ARM_JOINT_NAMES) or len(upper) != len(_ARM_JOINT_NAMES): + raise ValueError("Arm workspace bounds must contain one value per controlled joint.") + if any(not math.isfinite(value) for value in (*lower, *upper)): + raise ValueError("Arm workspace bounds must be finite.") + if any(low >= high for low, high in zip(lower, upper, strict=True)): + raise ValueError("Every arm workspace lower bound must be less than its upper bound.") + + +def _collision_properties(contact_margin: float = 0.0, mujoco_priority: int = 0) -> list[CollisionFragment]: + """Build explicit Newton and MuJoCo contact properties for one collider.""" + return [ + UsdPhysicsCollisionCfg(collision_enabled=True), + NewtonCollisionCfg(contact_margin=contact_margin, contact_gap=_CONTACT_GAP), + MujocoCollisionCfg( + condim=3, + priority=mujoco_priority, + solimp=_MUJOCO_SOLIMP, + solmix=1.0, + solref=_MUJOCO_SOLREF, + ), + ] + + +def _srgb_to_linear_channel(value: float) -> float: + """Convert an sRGB channel to the linear value expected by USD displayColor.""" + if value <= 0.04045: + return value / 12.92 + return ((value + 0.055) / 1.055) ** 2.4 + + +@configclass +class ActionsCfg: + """Relative arm and binary gripper actions.""" + + arm_action = mdp.ConveyorRelativeJointPositionActionCfg( + asset_name="robot", + joint_names=list(_ARM_JOINT_NAMES), + preserve_order=True, + scale=0.12, + max_delta=0.12, + ) + gripper_action = mdp.ResetBufferedGripperActionCfg( + asset_name="robot", + joint_names=list(_FINGER_JOINT_NAMES), + open_command_expr={"panda_finger_joint.*": 0.04}, + close_command_expr={"panda_finger_joint.*": 0.0}, + force_close_steps=5, + ) + + +@configclass +class CommandsCfg: + """Success-driven cube-transfer command.""" + + transfer = mdp.ConveyorTransferCommandCfg( + reset_event_name="reset_from_state_table", + minimum_subgoal_steps=2, + hold_steps=3, + lateral_tolerance=0.055, + maximum_cube_speed=0.65, + minimum_finger_position=0.027, + minimum_tool_clearance=0.055, + minimum_progress_steps=3, + minimum_progress=0.35, + maximum_target_potential=5.0, + minimum_acquisition_lift=0.025, + maximum_acquisition_tool_distance=0.075, + maximum_acquisition_finger_position=0.030, + ) + + +@configclass +class ObservationsCfg: + """Policy observations with stable cube identity and transfer commands.""" + + @configclass + class PolicyCfg(ObsGroup): + """Fully observed transfer policy input.""" + + joint_pos = ObsTerm( + func=mdp.joint_pos_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=list(_ARM_JOINT_NAMES), preserve_order=True)}, + ) + joint_vel = ObsTerm( + func=mdp.joint_vel_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=list(_ARM_JOINT_NAMES), preserve_order=True)}, + ) + gripper_pos = ObsTerm( + func=mdp.gripper_joint_positions, + params={"robot_cfg": SceneEntityCfg("robot", joint_names=list(_FINGER_JOINT_NAMES), preserve_order=True)}, + ) + objects = ObsTerm(func=mdp.transfer_object_observation) + active_transfer = ObsTerm(func=mdp.active_transfer_features) + target_cube = ObsTerm(func=mdp.target_cube_one_hot) + cube_conveyors = ObsTerm(func=mdp.cube_conveyor_state) + target_side = ObsTerm(func=mdp.target_side_one_hot) + eef_velocity = ObsTerm(func=mdp.end_effector_velocity) + eef_axes = ObsTerm(func=mdp.end_effector_axes) + last_action = ObsTerm(func=mdp.last_action) + + def __post_init__(self) -> None: + self.enable_corruption = False + self.concatenate_terms = True + + policy: PolicyCfg = PolicyCfg() + + +@configclass +class EventCfg: + """Restore validated physical reset states.""" + + reset_all = EventTerm(func=mdp.reset_scene_to_default, mode="reset") + reset_from_state_table = EventTerm( + func=mdp.ConveyorResetStateTable, + mode="reset", + params={ + "fixed_recipe": None, + "fixed_variant_id": None, + "fixed_target_cube_id": None, + "fixed_source_side_id": None, + "belt_start_x_range": (0.30, 0.82), + "cube_position_noise": 0.015, + "arm_joint_noise": 0.015, + }, + ) + + +@configclass +class RewardsCfg: + """Sparse transfer completion plus safety regularization rewards.""" + + success = RewTerm( + func=mdp.transfer_success_reward, + params={"command_name": "transfer"}, + weight=600.0, + ) + failure = RewTerm( + func=mdp.terminal_failure, + params={"command_name": "transfer"}, + weight=-60.0, + ) + arm_action_l2 = RewTerm( + func=mdp.action_term_l2, + params={"action_name": "arm_action"}, + weight=-1.0e-3, + ) + action_rate_l2 = RewTerm( + func=mdp.finite_action_rate_l2, + params={"action_names": ("arm_action", "gripper_action")}, + weight=-1.0e-3, + ) + joint_velocity_l2 = RewTerm( + func=mdp.finite_joint_velocity_l2, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=list(_ARM_JOINT_NAMES), preserve_order=True)}, + weight=-1.0e-4, + ) + + +@configclass +class TerminationsCfg: + """Safety failures and bounded training sequences.""" + + cube_out_of_workspace = DoneTerm(func=mdp.cube_out_of_workspace) + invalid_action = DoneTerm( + func=invalid_policy_action, + params={"action_names": ("arm_action", "gripper_action")}, + ) + nonfinite_scene_state = DoneTerm(func=mdp.nonfinite_scene_state) + subgoal_time_out = DoneTerm( + func=mdp.subgoal_time_out, + params={"timeout_s": _SUBGOAL_TIMEOUT_S, "command_name": "transfer"}, + time_out=True, + ) + transfer_sequence_time_out = DoneTerm( + func=mdp.transfer_sequence_time_out, + params={"maximum_transfers": _TRANSFER_SEQUENCE_LENGTH, "command_name": "transfer"}, + time_out=True, + ) + + +@configclass +class CurriculumCfg: + """Adaptive phase-balanced reset-state sampling.""" + + reset_sampling = CurrTerm( + func=mdp.ConveyorResetCurriculum, + params={ + "command_name": "transfer", + # Shared target-rate monitor keeps each physical reset row near the + # policy's 50% competence frontier without stale early outcomes. + "success_monitor": mdp.SuccessMonitorCfg( + monitored_history_len=50, + target_success_rate=0.5, + kappa=1.0, + temperature=1.0, + ), + # Keep a deployment-facing stream while the remaining starts + # adapt around the rolling pickup-to-placement frontier. Every + # recipe, cube identity, and direction retains equal total mass. + "deployment_probability_initial": 0.35, + "deployment_probability_final": 0.90, + "deployment_progress_start": 0.45, + "deployment_progress_end": 0.80, + "deployment_coverage_target": 0.50, + # Optional staged-training control. None keeps the deployable + # bidirectional task; a side id can focus the same adaptive reset + # distribution on one weak direction without changing reset rows. + "fixed_source_side_id": None, + }, + ) + + +@configclass +class ConveyorForceCfg: + """Configuration for force-based conveyor traction.""" + + speed: float = 0.35 + """Tangential conveyor surface speed [m/s].""" + + friction: float = 0.5 + """Coulomb friction coefficient used to limit traction.""" + + normal_threshold: float = 0.997 + """Minimum upward contact-normal alignment in the range [0, 1].""" + + startup_duration_s: float = 1.0 + """Duration over which conveyor traction ramps to full speed [s].""" + + transported_body_pattern: str = r"(?:^|/)Cube_?[0-3](?:/|$)" + """Regular expression selecting rigid bodies that receive conveyor forces.""" + + transported_body_count_per_env: int = 4 + """Expected number of transported rigid bodies in each environment.""" + + def __post_init__(self) -> None: + """Validate conveyor force parameters.""" + self.validate_config() + + def validate_config(self) -> None: + """Validate final conveyor-force values after overrides are applied.""" + if not math.isfinite(self.speed) or self.speed < 0.0: + raise ValueError(f"Conveyor speed must be non-negative, got {self.speed}.") + if not math.isfinite(self.friction) or self.friction < 0.0: + raise ValueError(f"Conveyor friction must be non-negative, got {self.friction}.") + if not math.isfinite(self.normal_threshold) or not 0.0 <= self.normal_threshold <= 1.0: + raise ValueError(f"Conveyor normal threshold must be in [0, 1], got {self.normal_threshold}.") + if not math.isfinite(self.startup_duration_s) or self.startup_duration_s <= 0.0: + raise ValueError(f"Conveyor startup duration must be positive, got {self.startup_duration_s}.") + if self.transported_body_count_per_env <= 0: + raise ValueError( + "Conveyor transported body count per environment must be positive, got " + f"{self.transported_body_count_per_env}." + ) + try: + re.compile(self.transported_body_pattern) + except re.error as exc: + raise ValueError(f"Invalid conveyor transported-body pattern: {self.transported_body_pattern!r}.") from exc + + +@sim_utils.clone +def _spawn_shape_with_display_color( + prim_path: str, + cfg: sim_utils.ShapeCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +): + """Spawn a primitive and author a renderer-independent USD display color.""" + if isinstance(cfg, sim_utils.CuboidCfg): + prim = sim_utils.spawn_cuboid(prim_path, cfg, translation, orientation, **kwargs) + elif isinstance(cfg, sim_utils.CylinderCfg): + prim = sim_utils.spawn_cylinder(prim_path, cfg, translation, orientation, **kwargs) + elif isinstance(cfg, sim_utils.CapsuleCfg): + prim = sim_utils.spawn_capsule(prim_path, cfg, translation, orientation, **kwargs) + elif isinstance(cfg, sim_utils.SphereCfg): + prim = sim_utils.spawn_sphere(prim_path, cfg, translation, orientation, **kwargs) + else: + raise TypeError(f"Unsupported colored primitive configuration: {type(cfg).__name__}") + + if cfg.visual_material is not None: + from pxr import Usd, UsdGeom + + display_color = tuple(_srgb_to_linear_channel(value) for value in cfg.visual_material.diffuse_color) + for child in Usd.PrimRange(prim): + if child.IsA(UsdGeom.Gprim): + UsdGeom.Gprim(child).CreateDisplayColorAttr([display_color]) + return prim + + +@sim_utils.clone +def _spawn_hidden_collision_mesh( + prim_path: str, + cfg: sim_utils.MeshCustomCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +): + """Spawn a collision-only custom mesh and hide it from all visualizers.""" + prim = sim_utils.spawn_mesh_custom(prim_path, cfg, translation, orientation, **kwargs) + sim_utils.set_prim_visibility(prim, False) + return prim + + +def _static_cuboid( + prim_path: str, + size: tuple[float, float, float], + pos: tuple[float, float, float], + color: tuple[float, float, float], + rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), + friction: float = 0.7, + roughness: float = 0.75, + metallic: float = 0.0, +) -> AssetBaseCfg: + """Build a static colliding cuboid configuration.""" + spawn = sim_utils.CuboidCfg( + size=size, + collision_props=_collision_properties(), + physics_material=RigidBodyMaterialBaseCfg( + static_friction=friction, + dynamic_friction=friction, + restitution=0.0, + ), + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=color, + roughness=roughness, + metallic=metallic, + ), + ) + spawn.func = _spawn_shape_with_display_color + return AssetBaseCfg( + prim_path=prim_path, + init_state=AssetBaseCfg.InitialStateCfg(pos=pos, rot=rot), + spawn=spawn, + ) + + +def _visual_mesh( + prim_path: str, + spec: MeshSpec, + color: tuple[float, float, float], + roughness: float, + metallic: float, +) -> AssetBaseCfg: + """Build a non-colliding custom mesh used only for rendering.""" + return AssetBaseCfg( + prim_path=prim_path, + spawn=sim_utils.MeshCustomCfg( + vertices=spec.vertices, + faces=spec.faces, + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=color, + roughness=roughness, + metallic=metallic, + ), + ), + ) + + +def _collision_material(friction: float) -> NewtonMaterialPropertiesCfg: + """Build a contact material for surfaces using the task-wide raw MuJoCo response.""" + return NewtonMaterialPropertiesCfg( + static_friction=friction, + dynamic_friction=friction, + restitution=0.0, + ) + + +def _hidden_collision_mesh( + prim_path: str, + spec: MeshSpec, + friction: float, + mujoco_priority: int, +) -> AssetBaseCfg: + """Build a hidden static triangle-mesh collider.""" + spawn = sim_utils.MeshCustomCfg( + vertices=spec.vertices, + faces=spec.faces, + visible=False, + collision_props=_collision_properties(mujoco_priority=mujoco_priority), + physics_material=_collision_material(friction), + ) + spawn.func = _spawn_hidden_collision_mesh + return AssetBaseCfg(prim_path=prim_path, spawn=spawn) + + +def _hidden_collision_cuboid( + prim_path: str, + spec: CuboidSpec, + friction: float, + mujoco_priority: int, +) -> AssetBaseCfg: + """Build a hidden native cuboid collider.""" + return AssetBaseCfg( + prim_path=prim_path, + init_state=AssetBaseCfg.InitialStateCfg(pos=spec.position), + spawn=sim_utils.CuboidCfg( + size=spec.size, + visible=False, + collision_props=_collision_properties(mujoco_priority=mujoco_priority), + physics_material=_collision_material(friction), + ), + ) + + +def _hidden_collision_geometry( + prim_path: str, + spec: MeshSpec | CuboidSpec, + friction: float, + mujoco_priority: int, +) -> AssetBaseCfg: + """Build hidden collision geometry while preferring native primitives where possible.""" + if isinstance(spec, CuboidSpec): + return _hidden_collision_cuboid( + prim_path=prim_path, + spec=spec, + friction=friction, + mujoco_priority=mujoco_priority, + ) + return _hidden_collision_mesh( + prim_path=prim_path, + spec=spec, + friction=friction, + mujoco_priority=mujoco_priority, + ) + + +def _cube( + name: str, + color: tuple[float, float, float], + pos: tuple[float, float, float], +) -> RigidObjectCfg: + """Build one numbered dynamic transfer cube.""" + spawn = sim_utils.CuboidCfg( + size=(0.04, 0.04, 0.04), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=color, roughness=0.75), + ) + spawn.rigid_props = _DYNAMIC_PROPERTIES + spawn.mass_props = sim_utils.MassPropertiesCfg(mass=0.05) + spawn.collision_props = _collision_properties(contact_margin=_CUBE_CONTACT_MARGIN) + spawn.physics_material = NewtonMaterialPropertiesCfg( + # The belt's higher MuJoCo contact priority overrides this friction + # only for belt/cube pairs, leaving physical finger/cube friction. + static_friction=0.8, + dynamic_friction=0.6, + restitution=0.0, + ) + spawn.func = _spawn_shape_with_display_color + return RigidObjectCfg( + prim_path=f"{{ENV_REGEX_NS}}/{name}", + init_state=RigidObjectCfg.InitialStateCfg(pos=pos), + spawn=spawn, + ) + + +@configclass +class ConveyorFrankaSceneCfg(InteractiveSceneCfg): + """Scene with two counter-rotating racetrack conveyors around a table-mounted Franka.""" + + # Use the MuJoCo Menagerie-derived model with explicit manipulation gains. + robot = FRANKA_PANDA_CONVEYOR_CFG.replace( + prim_path="{ENV_REGEX_NS}/Robot", + init_state=ArticulationCfg.InitialStateCfg( + joint_pos={ + "panda_joint1": 0.0, + "panda_joint2": -0.35, + "panda_joint3": 0.0, + "panda_joint4": -2.35, + "panda_joint5": 0.0, + "panda_joint6": 2.0, + "panda_joint7": 0.78, + "panda_finger_joint.*": 0.04, + } + ), + ) + + tabletop = _static_cuboid( + prim_path="{ENV_REGEX_NS}/Tabletop", + size=(2.0, 1.9, 0.08), + pos=(0.50, 0.0, -0.04), + color=(0.32, 0.34, 0.37), + ) + table_pedestal = _static_cuboid( + prim_path="{ENV_REGEX_NS}/TablePedestal", + size=(0.75, 0.55, 0.76), + pos=(0.25, 0.0, -0.46), + color=(0.18, 0.20, 0.23), + ) + + cube_0 = _cube("Cube0", CUBE_COLORS[0], (CUBE_INNER_SLOT_X, BELT_INNER_STRAIGHT_Y, 0.06)) + cube_1 = _cube("Cube1", CUBE_COLORS[1], (CUBE_OUTER_SLOT_X, BELT_OUTER_STRAIGHT_Y, 0.06)) + cube_2 = _cube("Cube2", CUBE_COLORS[2], (CUBE_INNER_SLOT_X, -BELT_INNER_STRAIGHT_Y, 0.06)) + cube_3 = _cube("Cube3", CUBE_COLORS[3], (CUBE_OUTER_SLOT_X, -BELT_OUTER_STRAIGHT_Y, 0.06)) + + ground = AssetBaseCfg( + prim_path="/World/GroundPlane", + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -0.85)), + spawn=sim_utils.GroundPlaneCfg(), + ) + dome_light = AssetBaseCfg( + prim_path="/World/DomeLight", + spawn=sim_utils.DomeLightCfg(color=(0.8, 0.8, 0.8), intensity=2500.0), + ) + + def __post_init__(self) -> None: + """Generate visual belts, velocity-field collision sections, and guardrails.""" + for side in ("Left", "Right"): + belt_spec = belt_mesh_spec(side) + setattr( + self, + f"conveyor_{side.lower()}_belt_visual", + _visual_mesh( + prim_path=f"{{ENV_REGEX_NS}}/{belt_spec.name}", + spec=belt_spec, + color=BELT_COLOR, + roughness=0.9, + metallic=0.0, + ), + ) + + section_keys = ("top_straight", "bottom_straight", "right_turn", "left_turn") + for section_key, spec in zip(section_keys, belt_collision_geometry_specs(side), strict=True): + setattr( + self, + f"conveyor_{side.lower()}_{section_key}_collision", + _hidden_collision_geometry( + prim_path=f"{{ENV_REGEX_NS}}/{spec.name}", + spec=spec, + # MuJoCo requires a tiny positive value even though the force driver, + # rather than solver friction, supplies the belt motion. + friction=1.1e-5, + # Override cube friction only for collision-section/cube pairs. + mujoco_priority=1, + ), + ) + + for spec in guard_mesh_specs(side): + boundary = "inner" if spec.name.endswith("Inner") else "outer" + setattr( + self, + f"guard_{side.lower()}_{boundary}_visual", + _visual_mesh( + prim_path=f"{{ENV_REGEX_NS}}/{spec.name}Visual", + spec=spec, + color=GUARD_COLOR, + roughness=0.3, + metallic=0.8, + ), + ) + setattr( + self, + f"guard_{side.lower()}_{boundary}_collision", + _hidden_collision_mesh( + prim_path=f"{{ENV_REGEX_NS}}/{spec.name}Collision", + spec=spec, + # The compact turns need freely sliding guide contacts; + # tangential rail friction can wedge a cube against the + # wall even though its belt drive remains valid. + friction=1.1e-5, + # Override the cube's grasp friction only for rail contacts. + mujoco_priority=1, + ), + ) + + +@configclass +class ConveyorFrankaEnvCfg(ManagerBasedRLEnvCfg): + """Manager-based RL task for commanded conveyor-to-conveyor cube transfer.""" + + scene: ConveyorFrankaSceneCfg = ConveyorFrankaSceneCfg(num_envs=256, env_spacing=3.0, replicate_physics=True) + conveyor_force: ConveyorForceCfg = ConveyorForceCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + observations: ObservationsCfg = ObservationsCfg() + events: EventCfg = EventCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + curriculum: CurriculumCfg = CurriculumCfg() + decimation: int = 2 + episode_length_s: float = _SUBGOAL_TIMEOUT_S * _TRANSFER_SEQUENCE_LENGTH + + sim: SimulationCfg = SimulationCfg( + dt=1.0 / 120.0, + render_interval=2, + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + solver="newton", + integrator="implicitfast", + # Per-environment capacities include headroom for all four cubes, + # the gripper, table, belts, and guard contacts without allocating + # the previous 256k-contact scene-wide driver buffers. + njmax=300, + nconmax=200, + impratio=1.0, + cone="elliptic", + update_data_interval=1, + iterations=100, + ls_iterations=50, + use_mujoco_contacts=False, + ccd_iterations=35, + ), + collision_cfg=NewtonCollisionPipelineCfg(), + # Manager decimation supplies the reference's two 120 Hz solves + # per 60 Hz policy step, with contacts refreshed before each solve. + collision_decimation=0, + default_shape_cfg=NewtonShapeCfg(margin=0.0, gap=_CONTACT_GAP, ke=2.5e3, kd=100.0), + num_substeps=1, + use_cuda_graph=True, + # Import render-only geometry only when a visualizer or camera needs it. + load_visual_shapes=None, + ), + use_newton_actuators=True, + ) + + def __post_init__(self) -> None: + # Any visualizer selected at runtime receives these shared camera hints. + # Newton camera pose: position (2.13, 0.0, 1.0), pitch -23.9 degrees, + # yaw 180 degrees. The look-at point is one unit along that view ray. + self.sim.default_visualizer_cfg = VisualizerCfg( + eye=(2.13, 0.0, 1.0), + lookat=(1.2157460448, 0.0, 0.5948584132), + max_visible_envs=1, + randomly_sample_visible_envs=False, + ) + + def validate_config(self) -> None: + """Validate the final task configuration after command-line overrides are applied.""" + _validate_common_config(self) + physics = self.sim.physics + if not isinstance(physics, NewtonCfg) or not isinstance(physics.solver_cfg, MJWarpSolverCfg): + raise ValueError("The conveyor force driver requires the Newton MJWarp backend.") + if physics.solver_cfg.use_mujoco_contacts: + raise ValueError("The conveyor force driver requires the Newton collision-pipeline contact path.") + if physics.collision_cfg is None: + raise ValueError("The conveyor force driver requires an explicit Newton collision pipeline.") + + def play_mode(self) -> None: + """Run continuing transfers from evenly distributed moving-belt starts.""" + super().play_mode() + self.scene.num_envs = min(self.scene.num_envs, 8) + self.events.reset_from_state_table.params["fixed_recipe"] = int(mdp.ConveyorResetRecipe.BELT) + self.events.reset_from_state_table.params["fixed_variant_id"] = mdp.BELT_DEPLOYMENT_VARIANT + self.events.reset_from_state_table.params["cube_position_noise"] = 0.0 + # Successful placements already transition to a new commanded cube. + # Playback removes training-only refreshes and runs until physics leaves + # the recoverable workspace. + self.terminations.subgoal_time_out = None + self.terminations.transfer_sequence_time_out = None + self.curriculum = None diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_physx_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_physx_env_cfg.py new file mode 100644 index 000000000000..122f9b374bc2 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_physx_env_cfg.py @@ -0,0 +1,421 @@ +# 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 + +"""CPU-only native-PhysX configuration for conveyor-Franka policy playback.""" + +from __future__ import annotations + +import functools +from dataclasses import replace + +from isaaclab_physx.physics import PhysxCfg, apply_surface_velocity_api +from isaaclab_physx.sim.schemas import PhysxCollisionCfg, PhysxSDFMeshCfg +from isaaclab_physx.sim.spawners.materials import PhysxRigidBodyMaterialCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from isaaclab.physics import SurfaceVelocitySpec +from isaaclab.sim import SimulationCfg +from isaaclab.sim.schemas import CollisionFragment, UsdPhysicsCollisionCfg +from isaaclab.utils.configclass import configclass + +from .conveyor_franka_env_cfg import ( + _CONTACT_GAP, + _CUBE_CONTACT_MARGIN, + ConveyorFrankaEnvCfg, + ConveyorFrankaSceneCfg, + _spawn_hidden_collision_mesh, + _spawn_shape_with_display_color, + _validate_common_config, + _visual_mesh, +) +from .conveyor_geometry import ( + BELT_COLOR, + BELT_INNER_STRAIGHT_Y, + BELT_OUTER_STRAIGHT_Y, + CUBE_COLORS, + CUBE_INNER_SLOT_X, + CUBE_OUTER_SLOT_X, + GUARD_COLOR, + ConveyorSectionSpec, + CuboidSpec, + MeshSpec, + belt_collision_section_specs, + belt_mesh_spec, + guard_mesh_specs, +) +from .franka_robot_cfg import FRANKA_PANDA_CONVEYOR_PHYSX_CFG + +_PHYSX_DYNAMIC_PROPERTIES = sim_utils.RigidBodyBaseCfg() +_PHYSX_KINEMATIC_PROPERTIES = sim_utils.RigidBodyBaseCfg( + rigid_body_enabled=True, + kinematic_enabled=True, + disable_gravity=True, +) + + +def _physx_collision_properties(contact_offset: float = 0.005) -> list[CollisionFragment]: + """Build the standard collision and PhysX offset fragments for one collider.""" + return [ + UsdPhysicsCollisionCfg(collision_enabled=True), + PhysxCollisionCfg(contact_offset=contact_offset, rest_offset=0.0), + ] + + +def _physx_material(friction: float) -> PhysxRigidBodyMaterialCfg: + """Build a deterministic PhysX material for a conveyor-task collision surface.""" + return PhysxRigidBodyMaterialCfg( + static_friction=friction, + dynamic_friction=friction, + restitution=0.0, + friction_combine_mode="min", + restitution_combine_mode="min", + ) + + +def _physx_static_cuboid( + prim_path: str, + size: tuple[float, float, float], + pos: tuple[float, float, float], + color: tuple[float, float, float], +) -> AssetBaseCfg: + """Build one static PhysX support cuboid.""" + spawn = sim_utils.CuboidCfg( + size=size, + collision_props=_physx_collision_properties(), + physics_material=_physx_material(0.7), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=color, roughness=0.75), + ) + spawn.func = _spawn_shape_with_display_color + return AssetBaseCfg( + prim_path=prim_path, + init_state=AssetBaseCfg.InitialStateCfg(pos=pos), + spawn=spawn, + ) + + +def _physx_cube( + name: str, + color: tuple[float, float, float], + pos: tuple[float, float, float], +) -> RigidObjectCfg: + """Build one numbered dynamic cube with native PhysX contact properties.""" + spawn = sim_utils.CuboidCfg( + size=(0.04, 0.04, 0.04), + rigid_props=_PHYSX_DYNAMIC_PROPERTIES, + mass_props=sim_utils.MassPropertiesCfg(mass=0.05), + collision_props=_physx_collision_properties(contact_offset=_CUBE_CONTACT_MARGIN), + physics_material=PhysxRigidBodyMaterialCfg( + static_friction=0.8, + dynamic_friction=0.6, + restitution=0.0, + restitution_combine_mode="min", + ), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=color, roughness=0.75), + ) + spawn.func = _spawn_shape_with_display_color + return RigidObjectCfg( + prim_path=f"{{ENV_REGEX_NS}}/{name}", + init_state=RigidObjectCfg.InitialStateCfg(pos=pos), + spawn=spawn, + ) + + +def _pivot_centered_section(section: ConveyorSectionSpec) -> tuple[ConveyorSectionSpec, tuple[float, float, float]]: + """Express a curved section around its rigid-body origin for native angular surface velocity.""" + if not section.belt.curved: + geometry = section.geometry + if not isinstance(geometry, CuboidSpec): + raise TypeError("Straight PhysX conveyor sections must use native cuboids.") + return section, geometry.position + + geometry = section.geometry + if not isinstance(geometry, MeshSpec): + raise TypeError("Curved PhysX conveyor sections must use closed meshes.") + pivot = section.belt.pivot_point + local_geometry = MeshSpec( + name=geometry.name, + vertices=tuple( + (vertex[0] - pivot[0], vertex[1] - pivot[1], vertex[2] - pivot[2]) for vertex in geometry.vertices + ), + faces=geometry.faces, + ) + local_belt = replace(section.belt, pivot_point=(0.0, 0.0, 0.0)) + return ConveyorSectionSpec(geometry=local_geometry, belt=local_belt), pivot + + +def physx_belt_section_specs( + side: str, + *, + velocity: float = 0.0, + friction_coefficient: float = 0.5, + contact_threshold: float = 0.997, +) -> tuple[tuple[ConveyorSectionSpec, tuple[float, float, float]], ...]: + """Return pivot-centered sections sharing the exact runtime PhysX semantics.""" + return tuple( + _pivot_centered_section(section) + for section in belt_collision_section_specs( + side, + velocity=velocity, + friction_coefficient=friction_coefficient, + contact_threshold=contact_threshold, + ) + ) + + +@sim_utils.clone +def _spawn_physx_conveyor_mesh( + prim_path: str, + cfg: sim_utils.MeshCustomCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + *, + belt_spec: SurfaceVelocitySpec, + **kwargs, +): + """Spawn one hidden SDF turn and author native surface velocity before PhysX parsing.""" + prim = sim_utils.spawn_mesh_custom(prim_path, cfg, translation, orientation, **kwargs) + sim_utils.set_prim_visibility(prim, False) + apply_surface_velocity_api(prim, belt_spec, velocity_scale=0.0) + return prim + + +@sim_utils.clone +def _spawn_physx_conveyor_cuboid( + prim_path: str, + cfg: sim_utils.CuboidCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + *, + belt_spec: SurfaceVelocitySpec, + **kwargs, +): + """Spawn one hidden analytic straight and author native surface velocity before PhysX parsing.""" + prim = sim_utils.spawn_cuboid(prim_path, cfg, translation, orientation, **kwargs) + sim_utils.set_prim_visibility(prim, False) + apply_surface_velocity_api(prim, belt_spec, velocity_scale=0.0) + return prim + + +def _physx_conveyor_collision( + prim_path: str, + section: ConveyorSectionSpec, + root_position: tuple[float, float, float], + friction: float, +) -> AssetBaseCfg: + """Build one native-velocity belt section with analytic or watertight-SDF collision.""" + geometry = section.geometry + if isinstance(geometry, CuboidSpec): + spawn = sim_utils.CuboidCfg( + size=geometry.size, + visible=False, + rigid_props=_PHYSX_KINEMATIC_PROPERTIES, + collision_props=_physx_collision_properties(contact_offset=_CONTACT_GAP), + physics_material=_physx_material(friction), + ) + spawn.func = functools.partial(_spawn_physx_conveyor_cuboid, belt_spec=section.belt) + else: + spawn = sim_utils.MeshCustomCfg( + vertices=geometry.vertices, + faces=geometry.faces, + visible=False, + rigid_props=_PHYSX_KINEMATIC_PROPERTIES, + collision_props=[ + *_physx_collision_properties(contact_offset=_CONTACT_GAP), + PhysxSDFMeshCfg(sdf_resolution=128, sdf_subgrid_resolution=6), + ], + collision_approximation="sdf", + physics_material=_physx_material(friction), + ) + spawn.func = functools.partial(_spawn_physx_conveyor_mesh, belt_spec=section.belt) + return AssetBaseCfg( + prim_path=prim_path, + init_state=AssetBaseCfg.InitialStateCfg(pos=root_position), + spawn=spawn, + ) + + +def _physx_guard_collision(prim_path: str, spec: MeshSpec) -> AssetBaseCfg: + """Build one freely sliding static guide collider.""" + spawn = sim_utils.MeshCustomCfg( + vertices=spec.vertices, + faces=spec.faces, + visible=False, + collision_props=_physx_collision_properties(), + collision_approximation="none", + physics_material=_physx_material(1.1e-5), + ) + spawn.func = _spawn_hidden_collision_mesh + return AssetBaseCfg(prim_path=prim_path, spawn=spawn) + + +@configclass +class ConveyorFrankaPhysxSceneCfg(ConveyorFrankaSceneCfg): + """PhysX scene preserving the Newton task's names, layout, and tensor contracts.""" + + robot = FRANKA_PANDA_CONVEYOR_PHYSX_CFG.replace( + prim_path="{ENV_REGEX_NS}/Robot", + init_state=ArticulationCfg.InitialStateCfg( + joint_pos={ + "panda_joint1": 0.0, + "panda_joint2": -0.35, + "panda_joint3": 0.0, + "panda_joint4": -2.35, + "panda_joint5": 0.0, + "panda_joint6": 2.0, + "panda_joint7": 0.78, + "panda_finger_joint.*": 0.04, + } + ), + ) + + tabletop = _physx_static_cuboid( + prim_path="{ENV_REGEX_NS}/Tabletop", + size=(2.0, 1.9, 0.08), + pos=(0.50, 0.0, -0.04), + color=(0.32, 0.34, 0.37), + ) + table_pedestal = _physx_static_cuboid( + prim_path="{ENV_REGEX_NS}/TablePedestal", + size=(0.75, 0.55, 0.76), + pos=(0.25, 0.0, -0.46), + color=(0.18, 0.20, 0.23), + ) + + cube_0 = _physx_cube("Cube0", CUBE_COLORS[0], (CUBE_INNER_SLOT_X, BELT_INNER_STRAIGHT_Y, 0.06)) + cube_1 = _physx_cube("Cube1", CUBE_COLORS[1], (CUBE_OUTER_SLOT_X, BELT_OUTER_STRAIGHT_Y, 0.06)) + cube_2 = _physx_cube("Cube2", CUBE_COLORS[2], (CUBE_INNER_SLOT_X, -BELT_INNER_STRAIGHT_Y, 0.06)) + cube_3 = _physx_cube("Cube3", CUBE_COLORS[3], (CUBE_OUTER_SLOT_X, -BELT_OUTER_STRAIGHT_Y, 0.06)) + + def __post_init__(self) -> None: + """Generate visuals plus native PhysX belt and guide collision bodies.""" + for side in ("Left", "Right"): + visual = belt_mesh_spec(side) + setattr( + self, + f"conveyor_{side.lower()}_belt_visual", + _visual_mesh( + prim_path=f"{{ENV_REGEX_NS}}/{visual.name}", + spec=visual, + color=BELT_COLOR, + roughness=0.9, + metallic=0.0, + ), + ) + + section_keys = ("top_straight", "bottom_straight", "right_turn", "left_turn") + for section_key, (section, root_position) in zip(section_keys, physx_belt_section_specs(side), strict=True): + setattr( + self, + f"conveyor_{side.lower()}_{section_key}_collision", + _physx_conveyor_collision( + prim_path=f"{{ENV_REGEX_NS}}/{section.geometry.name}", + section=section, + root_position=root_position, + friction=0.5, + ), + ) + + for guard in guard_mesh_specs(side): + boundary = "inner" if guard.name.endswith("Inner") else "outer" + setattr( + self, + f"guard_{side.lower()}_{boundary}_visual", + _visual_mesh( + prim_path=f"{{ENV_REGEX_NS}}/{guard.name}Visual", + spec=guard, + color=GUARD_COLOR, + roughness=0.3, + metallic=0.8, + ), + ) + setattr( + self, + f"guard_{side.lower()}_{boundary}_collision", + _physx_guard_collision(f"{{ENV_REGEX_NS}}/{guard.name}Collision", guard), + ) + + def build_conveyor_belt_specs( + self, + *, + velocity: float, + friction_coefficient: float, + contact_threshold: float, + ) -> tuple[SurfaceVelocitySpec, ...]: + """Return the same pivot-local belt descriptions used by the PhysX spawners.""" + return tuple( + section.belt + for side in ("Left", "Right") + for section, _ in physx_belt_section_specs( + side, + velocity=velocity, + friction_coefficient=friction_coefficient, + contact_threshold=contact_threshold, + ) + ) + + def configure_conveyor(self, *, friction_coefficient: float) -> None: + """Propagate a final command-line friction override into every belt material before spawning.""" + for side in ("left", "right"): + for section_key in ("top_straight", "bottom_straight", "right_turn", "left_turn"): + asset = getattr(self, f"conveyor_{side}_{section_key}_collision") + material = asset.spawn.physics_material + material.static_friction = friction_coefficient + material.dynamic_friction = friction_coefficient + + +@configclass +class ConveyorFrankaPhysxEnvCfg(ConveyorFrankaEnvCfg): + """Checkpoint-compatible CPU reference using native PhysX surface velocity. + + The supported and tested native ``PhysxSurfaceVelocityAPI`` path is CPU-only. Use + :class:`ConveyorFrankaEnvCfg` for scalable GPU simulation with Newton. + """ + + scene: ConveyorFrankaPhysxSceneCfg = ConveyorFrankaPhysxSceneCfg( + # USD surface-velocity commands are authored on the host. One environment + # is the useful default for this CPU reference; callers may opt into a + # small replicated batch explicitly. + num_envs=1, + env_spacing=3.0, + replicate_physics=True, + ) + sim: SimulationCfg = SimulationCfg( + # The supported GPU contact-modification path drops contacts for shapes + # with PhysxSurfaceVelocityAPI enabled. CPU PhysX preserves the authored + # API and its belt contacts; fail validation on an accidental CUDA + # override instead of letting cubes tunnel silently. + device="cpu", + dt=1.0 / 120.0, + render_interval=2, + physics=PhysxCfg( + solver_type=1, + solve_articulation_contact_last=True, + max_position_iteration_count=64, + max_velocity_iteration_count=16, + bounce_threshold_velocity=0.2, + friction_offset_threshold=0.01, + friction_correlation_distance=0.00625, + enable_ccd=True, + ), + use_fabric=True, + use_newton_actuators=True, + ) + + def validate_config(self) -> None: + """Validate PhysX selection and the shared checkpoint-facing policy contract.""" + _validate_common_config(self) + if not isinstance(self.sim.physics, PhysxCfg): + raise ValueError("The PhysX conveyor task requires the native Isaac Sim PhysX backend.") + if self.sim.device != "cpu": + raise ValueError( + "The native PhysX conveyor is CPU-only because enabled PhysxSurfaceVelocityAPI shapes can lose " + "contacts under GPU dynamics. Run this task with '--device cpu'; use the Newton task for GPU " + "simulation." + ) + if self.scene.robot.spawn.joint_drive_props is not None: + raise ValueError("The PhysX Franka must not author MuJoCo-only joint-drive properties.") + if self.scene.robot.spawn.rigid_props.disable_gravity is not True: + raise ValueError("The PhysX Franka must preserve the trained gravity-compensated policy contract.") diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_warehouse_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_warehouse_env.py new file mode 100644 index 000000000000..5cec7184b25e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_franka_warehouse_env.py @@ -0,0 +1,153 @@ +# 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 + +"""Four-slot policy playback over a physical parcel pool and USD-authored warehouse.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.envs.common import VecEnvStepReturn + +from .conveyor_cube_pool import ConveyorCubePool +from .conveyor_franka_env import ConveyorFrankaEnv + +if TYPE_CHECKING: + from pxr import Usd + + from .conveyor_franka_asset_env_cfg import ConveyorFrankaA09A12EnvCfg + + +class ConveyorFrankaWarehouseEnv(ConveyorFrankaEnv): + """Play the four-cube checkpoint over an individually tracked workcell parcel pool. + + Kit renders the authored USD materials, lights, and animation. Lightweight + Newton viewers show a static approximation of the warehouse dressing. + """ + + def __init__(self, cfg: ConveyorFrankaA09A12EnvCfg, render_mode: str | None = None, **kwargs): + self._warehouse_animation: list[tuple[Usd.Attribute, Usd.Attribute]] = [] + self._warehouse_arm_joint_ids: list[int] | None = None + cfg.scene._configure_route_assets(cfg.commands.transfer.parcel_colors) + super().__init__(cfg, render_mode=render_mode, **kwargs) + if not any(viz.cfg.visualizer_type == "kit" for viz in self.sim.visualizers): + return + self.sim.set_setting("/app/viewport/grid/enabled", False) + + from pxr import Usd + + from .conveyor_franka_asset_env_cfg import _presentation_layer + + camera = self.sim.stage.GetPrimAtPath("/OmniverseKit_Persp") + if camera: + # Kit owns this camera in its session layer; weaker root-layer edits are ignored. + viewer = cfg.sim.default_visualizer_cfg + with Usd.EditContext(self.sim.stage, self.sim.stage.GetSessionLayer()): + self.sim.set_camera_view(viewer.eye, viewer.lookat) + camera.GetAttribute("focalLength").Set(viewer.focal_length) + self._warehouse_source = Usd.Stage.Open(_presentation_layer(cfg.scene.warehouse_visual.spawn.usd_path)) + self._warehouse_period = self._warehouse_source.GetEndTimeCode() + self._warehouse_fps = self._warehouse_source.GetTimeCodesPerSecond() + for env_path in self.scene.env_prim_paths: + group = self._warehouse_source.GetPrimAtPath("/Warehouse/Parcels") + for source in group.GetChildren(): + target = self.sim.stage.GetPrimAtPath(f"{env_path}/WarehouseVisual/Parcels/{source.GetName()}") + for name in ("xformOp:translate", "xformOp:rotateXYZ"): + self._warehouse_animation.append((source.GetAttribute(name), target.GetAttribute(name))) + self.sim.add_render_callback("conveyor_warehouse_animation", self._animate_warehouse) + + def load_managers(self) -> None: + """Create slot identities before command, reward, and observation terms inspect cubes.""" + assets = tuple(self.scene[f"cube_{i}"] for i in range(self.cfg.conveyor_force.transported_body_count_per_env)) + self.conveyor_cube_pool = ConveyorCubePool(assets, self.num_envs, self.device) + super().load_managers() + + @staticmethod + def _in_workcell(positions: torch.Tensor) -> torch.Tensor: + """Identify parcels within the trained controller's local manipulation region [m].""" + return ( + (positions[..., 0] > -0.25) + & (positions[..., 0] < 1.5) + & (positions[..., 1].abs() < 0.6) + & (positions[..., 2] < 0.4) + # The low merge passes beside the placement bend but is still remote transport. + & ((positions[..., 2] < 0.12) | (positions[..., 0] < 1.05)) + ) + + def _adapt_policy_cube_state(self, positions, quaternions, velocities): + """Represent remote inventory as waiting slots while keeping local manipulation states exact.""" + origins = self.scene.env_origins[:, None, :] + local = positions - origins + remote = ~self._in_workcell(local) + waiting = local.clone() + waiting[..., 0] = 0.14 + 0.88 * torch.arange(1, 5, device=self.device)[None, :] / 5 + waiting[..., 1] = torch.where(local[..., 1] >= 0, 0.75, -0.75) + waiting[..., 2] = 0.06 + upright = torch.zeros_like(quaternions) + upright[..., 3] = 1.0 + transport_velocity = torch.zeros_like(velocities) + transport_velocity[..., 0] = self.cfg.conveyor_force.speed + return ( + torch.where(remote[..., None], waiting + origins, positions), + torch.where(remote[..., None], upright, quaternions), + torch.where(remote[..., None], transport_velocity, velocities), + ) + + def step(self, action: torch.Tensor) -> VecEnvStepReturn: + """Park while waiting for a misplaced parcel; dispatch runs through the command manager.""" + command = self.command_manager.get_term("transfer") + robot = self.scene["robot"] + if self._warehouse_arm_joint_ids is None: + self._warehouse_arm_joint_ids = robot.find_joints( + self.cfg.actions.arm_action.joint_names, preserve_order=True + )[0] + joint_ids = self._warehouse_arm_joint_ids + parked = torch.zeros_like(action) + parked[:, :7] = ( + (robot.data.default_joint_pos.torch[:, joint_ids] - robot.data.joint_pos.torch[:, joint_ids]) + / self.cfg.actions.arm_action.scale + ).clamp(-0.25, 0.25) + # Keep invalid actions visible to the existing sanitization and termination terms. + use_policy = command.has_target | ~torch.isfinite(action).all(dim=1) + return super().step(torch.where(use_policy[:, None], action, parked)) + + def _animate_warehouse(self, _event) -> None: + """Sample USD motion using policy time, independently of render frame rate.""" + from pxr import Sdf + + time_code = (self.common_step_counter * self.step_dt * self._warehouse_fps) % self._warehouse_period + with Sdf.ChangeBlock(): + for source, target in self._warehouse_animation: + target.Set(source.Get(time_code)) + + def _reset_idx(self, env_ids) -> None: + """Shuffle a mixed batch across the authored feeds in selected environments.""" + from .conveyor_warehouse_geometry import warehouse_parcel_positions + + if isinstance(env_ids, slice): + env_ids = self.scene._ALL_INDICES[env_ids] + pool = self.conveyor_cube_pool + pool.reset(env_ids) + super()._reset_idx(env_ids) + positions = torch.tensor(warehouse_parcel_positions(), device=self.device) + if self.cfg.commands.transfer.randomize_arrivals: + assignments = torch.rand((len(env_ids), len(pool.assets)), device=self.device).argsort(dim=1) + else: + assignments = torch.arange(len(pool.assets), device=self.device).expand(len(env_ids), -1) + for cube_id, cube in enumerate(pool.assets): + pose = cube.data.root_pose_w.torch[env_ids].clone() + pose[:, :3] = positions[assignments[:, cube_id]] + self.scene.env_origins[env_ids] + pose[:, 3:] = pose.new_tensor((0.0, 0.0, 0.0, 1.0)) + cube.write_root_pose_to_sim_index(root_pose=pose, env_ids=env_ids) + cube.write_root_velocity_to_sim_index(root_velocity=pose.new_zeros((len(env_ids), 6)), env_ids=env_ids) + + def close(self) -> None: + """Remove the presentation callback before releasing the shared task resources.""" + if getattr(self, "sim", None) is not None: + self.sim.remove_render_callback("conveyor_warehouse_animation") + self._warehouse_animation.clear() + super().close() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_geometry.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_geometry.py new file mode 100644 index 000000000000..b218cde91949 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_geometry.py @@ -0,0 +1,397 @@ +# 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 + +"""Shared racetrack geometry and velocity-field descriptions.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +from isaaclab.physics import SurfaceVelocitySpec + +BELT_COLOR = (0.09, 0.09, 0.09) +"""Dark-rubber color used by Newton's conveyor example.""" + +GUARD_COLOR = (0.66, 0.69, 0.74) +"""Brushed-metal color used by Newton's conveyor example.""" + +PARCEL_COLOR = (0.72, 0.55, 0.35) +CUBE_COLORS = ( + (0.15, 0.35, 0.90), + (0.90, 0.20, 0.15), + (0.15, 0.75, 0.25), + PARCEL_COLOR, +) +"""Stable sRGB colors for the four numbered transfer cubes.""" +"""Cardboard color used by Newton's conveyor example.""" + +BELT_CENTER_X = 0.58 +BELT_CENTER_Y = 0.51 +BELT_TURN_RADIUS = 0.24 +BELT_HALF_STRAIGHT = 0.44 +BELT_WIDTH = 0.15 +BELT_THICKNESS = 0.04 +BELT_TOP_Z = 0.04 +TURN_SEGMENT_COUNT = 96 + +# Four deployment slots place one cube on each straight run of the two +# racetracks. The x coordinates are mirrored about the racetrack center, so +# the inner/outer pair on a belt is separated by exactly half a lap. +BELT_INNER_STRAIGHT_Y = BELT_CENTER_Y - BELT_TURN_RADIUS +BELT_OUTER_STRAIGHT_Y = BELT_CENTER_Y + BELT_TURN_RADIUS +CUBE_INNER_SLOT_X = BELT_CENTER_X - BELT_HALF_STRAIGHT / 3.0 +CUBE_OUTER_SLOT_X = BELT_CENTER_X + BELT_HALF_STRAIGHT / 3.0 + +# Collision surfaces extend underneath the rails and overlap at section seams. +# This keeps the dynamic parcels on a continuous +Z-facing surface without +# exposing the belt prism's vertical side faces to the contact solver. +BELT_COLLISION_OVERHANG = 0.02 +BELT_COLLISION_SEAM_OVERLAP = 0.004 + +GUARD_THICKNESS = 0.018 +# Keep the rails below the parcel tops so both lanes remain easy to read from +# the default oblique camera. +GUARD_HEIGHT = 0.02 +GUARD_BASE_OVERLAP = 0.005 + + +@dataclass(frozen=True) +class MeshSpec: + """Triangle mesh and semantic name for one static racetrack component.""" + + name: str + vertices: tuple[tuple[float, float, float], ...] + faces: tuple[tuple[int, int, int], ...] + + +@dataclass(frozen=True) +class CuboidSpec: + """Native cuboid and semantic name for one static racetrack component.""" + + name: str + size: tuple[float, float, float] + position: tuple[float, float, float] + + +@dataclass(frozen=True) +class ConveyorSectionSpec: + """Task geometry paired with its backend-neutral conveyor description.""" + + geometry: MeshSpec | CuboidSpec + belt: SurfaceVelocitySpec + + +def belt_direction(side: str) -> float: + """Return ``1`` for clockwise motion and ``-1`` for counter-clockwise motion.""" + if side == "Left": + return 1.0 + if side == "Right": + return -1.0 + raise ValueError(f"Unknown conveyor side: {side!r}.") + + +def _racetrack_centerline(center_y: float) -> tuple[tuple[float, float, float, float], ...]: + """Sample a clockwise racetrack centerline with an outward normal at every point.""" + left_x = BELT_CENTER_X - BELT_HALF_STRAIGHT + right_x = BELT_CENTER_X + BELT_HALF_STRAIGHT + radius = BELT_TURN_RADIUS + points: list[tuple[float, float, float, float]] = [ + (left_x, center_y + radius, 0.0, 1.0), + (right_x, center_y + radius, 0.0, 1.0), + ] + + for index in range(1, TURN_SEGMENT_COUNT + 1): + angle = 0.5 * math.pi - index * math.pi / TURN_SEGMENT_COUNT + normal_x = math.cos(angle) + normal_y = math.sin(angle) + points.append((right_x + radius * normal_x, center_y + radius * normal_y, normal_x, normal_y)) + + points.append((left_x, center_y - radius, 0.0, -1.0)) + for index in range(1, TURN_SEGMENT_COUNT): + angle = -0.5 * math.pi - index * math.pi / TURN_SEGMENT_COUNT + normal_x = math.cos(angle) + normal_y = math.sin(angle) + points.append((left_x + radius * normal_x, center_y + radius * normal_y, normal_x, normal_y)) + return tuple(points) + + +def _racetrack_prism_mesh( + name: str, + center_y: float, + lateral_offset: float, + width: float, + z_min: float, + z_max: float, +) -> MeshSpec: + """Build one closed prism following a racetrack centerline.""" + centerline = _racetrack_centerline(center_y) + half_width = 0.5 * width + outer_offset = lateral_offset + half_width + inner_offset = lateral_offset - half_width + + outer_top = [(x + nx * outer_offset, y + ny * outer_offset, z_max) for x, y, nx, ny in centerline] + inner_top = [(x + nx * inner_offset, y + ny * inner_offset, z_max) for x, y, nx, ny in centerline] + outer_bottom = [(x + nx * outer_offset, y + ny * outer_offset, z_min) for x, y, nx, ny in centerline] + inner_bottom = [(x + nx * inner_offset, y + ny * inner_offset, z_min) for x, y, nx, ny in centerline] + points = tuple(inner_top + outer_top + inner_bottom + outer_bottom) + + count = len(centerline) + outer_top_offset = count + inner_bottom_offset = 2 * count + outer_bottom_offset = 3 * count + indices: list[int] = [] + for index in range(count): + next_index = (index + 1) % count + inner_top_i = index + inner_top_j = next_index + outer_top_i = outer_top_offset + index + outer_top_j = outer_top_offset + next_index + inner_bottom_i = inner_bottom_offset + index + inner_bottom_j = inner_bottom_offset + next_index + outer_bottom_i = outer_bottom_offset + index + outer_bottom_j = outer_bottom_offset + next_index + + # Top, bottom, outer wall, and inner wall; two triangles per surface. + indices.extend((inner_top_i, outer_top_i, outer_top_j, inner_top_i, outer_top_j, inner_top_j)) + indices.extend( + ( + inner_bottom_i, + inner_bottom_j, + outer_bottom_j, + inner_bottom_i, + outer_bottom_j, + outer_bottom_i, + ) + ) + indices.extend( + ( + outer_bottom_i, + outer_bottom_j, + outer_top_j, + outer_bottom_i, + outer_top_j, + outer_top_i, + ) + ) + indices.extend( + ( + inner_bottom_i, + inner_top_i, + inner_top_j, + inner_bottom_i, + inner_top_j, + inner_bottom_j, + ) + ) + + # The centerline is sampled clockwise so its tangent matches the positive + # conveyor direction. The face pattern above is written for a + # counter-clockwise ring, so reverse every triangle to keep its collision + # normal outward (most importantly, the belt top must point +Z). + for triangle_start in range(0, len(indices), 3): + indices[triangle_start + 1], indices[triangle_start + 2] = ( + indices[triangle_start + 2], + indices[triangle_start + 1], + ) + + return MeshSpec( + name=name, + vertices=points, + faces=tuple(tuple(indices[offset : offset + 3]) for offset in range(0, len(indices), 3)), + ) + + +def belt_mesh_spec(side: str) -> MeshSpec: + """Build the seamless, watertight visual mesh for one conveyor belt.""" + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + return _racetrack_prism_mesh( + name=f"Conveyor{side}BeltVisual", + center_y=center_y, + lateral_offset=0.0, + width=BELT_WIDTH, + z_min=BELT_TOP_Z - BELT_THICKNESS, + z_max=BELT_TOP_Z, + ) + + +def _straight_collision_cuboid(name: str, center_y: float) -> CuboidSpec: + """Build one solid native cuboid for a straight conveyor section.""" + half_width = 0.5 * BELT_WIDTH + BELT_COLLISION_OVERHANG + x_min = BELT_CENTER_X - BELT_HALF_STRAIGHT - BELT_COLLISION_SEAM_OVERLAP + x_max = BELT_CENTER_X + BELT_HALF_STRAIGHT + BELT_COLLISION_SEAM_OVERLAP + return CuboidSpec( + name=name, + size=(x_max - x_min, 2.0 * half_width, BELT_THICKNESS), + position=(0.5 * (x_min + x_max), center_y, BELT_TOP_Z - 0.5 * BELT_THICKNESS), + ) + + +def _turn_collision_mesh( + name: str, pivot_x: float, center_y: float, start_angle: float, angle_span: float = math.pi +) -> MeshSpec: + """Build a closed annular turn with a +Z top surface and an angular span [rad].""" + half_width = 0.5 * BELT_WIDTH + BELT_COLLISION_OVERHANG + inner_radius = BELT_TURN_RADIUS - half_width + outer_radius = BELT_TURN_RADIUS + half_width + angle_overlap = BELT_COLLISION_SEAM_OVERLAP / BELT_TURN_RADIUS + angle_start = start_angle - angle_overlap + segment_count = round(TURN_SEGMENT_COUNT * angle_span / math.pi) + angle_step = (angle_span + 2.0 * angle_overlap) / segment_count + angles = tuple(angle_start + index * angle_step for index in range(segment_count + 1)) + + inner_top = tuple( + (pivot_x + inner_radius * math.cos(angle), center_y + inner_radius * math.sin(angle), BELT_TOP_Z) + for angle in angles + ) + outer_top = tuple( + (pivot_x + outer_radius * math.cos(angle), center_y + outer_radius * math.sin(angle), BELT_TOP_Z) + for angle in angles + ) + bottom_z = BELT_TOP_Z - BELT_THICKNESS + inner_bottom = tuple((x, y, bottom_z) for x, y, _ in inner_top) + outer_bottom = tuple((x, y, bottom_z) for x, y, _ in outer_top) + + count = len(inner_top) + outer_top_offset = count + inner_bottom_offset = 2 * count + outer_bottom_offset = 3 * count + faces: list[tuple[int, int, int]] = [] + for index in range(segment_count): + next_index = index + 1 + inner_top_i = index + inner_top_j = next_index + outer_top_i = outer_top_offset + index + outer_top_j = outer_top_offset + next_index + inner_bottom_i = inner_bottom_offset + index + inner_bottom_j = inner_bottom_offset + next_index + outer_bottom_i = outer_bottom_offset + index + outer_bottom_j = outer_bottom_offset + next_index + faces.extend( + ( + # Top, bottom, outer wall, and inner wall. + (inner_top_i, outer_top_i, outer_top_j), + (inner_top_i, outer_top_j, inner_top_j), + (inner_bottom_i, inner_bottom_j, outer_bottom_j), + (inner_bottom_i, outer_bottom_j, outer_bottom_i), + (outer_bottom_i, outer_bottom_j, outer_top_j), + (outer_bottom_i, outer_top_j, outer_top_i), + (inner_bottom_i, inner_top_i, inner_top_j), + (inner_bottom_i, inner_top_j, inner_bottom_j), + ) + ) + + # Close both radial ends of the annular prism. + end = segment_count + faces.extend( + ( + (0, inner_bottom_offset, outer_bottom_offset), + (0, outer_bottom_offset, outer_top_offset), + (end, outer_top_offset + end, outer_bottom_offset + end), + (end, outer_bottom_offset + end, inner_bottom_offset + end), + ) + ) + return MeshSpec( + name=name, + vertices=inner_top + outer_top + inner_bottom + outer_bottom, + faces=tuple(faces), + ) + + +def belt_collision_geometry_specs(side: str) -> tuple[CuboidSpec | MeshSpec, ...]: + """Build native straight and closed-mesh turn collision geometry for one racetrack.""" + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + left_x = BELT_CENTER_X - BELT_HALF_STRAIGHT + right_x = BELT_CENTER_X + BELT_HALF_STRAIGHT + return ( + _straight_collision_cuboid(f"Conveyor{side}TopStraightCollision", center_y + BELT_TURN_RADIUS), + _straight_collision_cuboid(f"Conveyor{side}BottomStraightCollision", center_y - BELT_TURN_RADIUS), + _turn_collision_mesh(f"Conveyor{side}RightTurnCollision", right_x, center_y, -0.5 * math.pi), + _turn_collision_mesh(f"Conveyor{side}LeftTurnCollision", left_x, center_y, 0.5 * math.pi), + ) + + +def belt_collision_section_specs( + side: str, + *, + velocity: float = 0.0, + friction_coefficient: float = 0.7, + contact_threshold: float = 0.997, + enabled: bool = True, +) -> tuple[ConveyorSectionSpec, ConveyorSectionSpec, ConveyorSectionSpec, ConveyorSectionSpec]: + """Build collision geometry and authored conveyor intent for one racetrack.""" + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + left_x = BELT_CENTER_X - BELT_HALF_STRAIGHT + right_x = BELT_CENTER_X + BELT_HALF_STRAIGHT + direction = belt_direction(side) + top_straight, bottom_straight, right_turn, left_turn = belt_collision_geometry_specs(side) + + def belt( + geometry: MeshSpec | CuboidSpec, + travel_direction: tuple[float, float, float], + *, + curved: bool = False, + pivot_point: tuple[float, float, float] = (0.0, 0.0, 0.0), + radius: float | None = None, + ) -> ConveyorSectionSpec: + return ConveyorSectionSpec( + geometry=geometry, + belt=SurfaceVelocitySpec( + prim_path=f"{{ENV_REGEX_NS}}/{geometry.name}", + velocity=velocity, + enabled=enabled, + direction=travel_direction, + curved=curved, + pivot_point=pivot_point, + radius=radius, + contact_threshold=contact_threshold, + friction_coefficient=friction_coefficient, + ), + ) + + return ( + belt(top_straight, (direction, 0.0, 0.0)), + belt(bottom_straight, (-direction, 0.0, 0.0)), + belt( + right_turn, + (0.0, 0.0, -direction), + curved=True, + pivot_point=(right_x, center_y, 0.0), + radius=BELT_TURN_RADIUS, + ), + belt( + left_turn, + (0.0, 0.0, -direction), + curved=True, + pivot_point=(left_x, center_y, 0.0), + radius=BELT_TURN_RADIUS, + ), + ) + + +def guard_mesh_specs(side: str) -> tuple[MeshSpec, MeshSpec]: + """Build seamless inner and outer guardrail meshes for one racetrack.""" + center_y = BELT_CENTER_Y if side == "Left" else -BELT_CENTER_Y + rail_offset = 0.5 * (BELT_WIDTH + GUARD_THICKNESS) + z_min = BELT_TOP_Z - GUARD_BASE_OVERLAP + z_max = BELT_TOP_Z + GUARD_HEIGHT + return ( + _racetrack_prism_mesh( + name=f"Guard{side}Inner", + center_y=center_y, + lateral_offset=-rail_offset, + width=GUARD_THICKNESS, + z_min=z_min, + z_max=z_max, + ), + _racetrack_prism_mesh( + name=f"Guard{side}Outer", + center_y=center_y, + lateral_offset=rail_offset, + width=GUARD_THICKNESS, + z_min=z_min, + z_max=z_max, + ), + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_goal_selector.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_goal_selector.py new file mode 100644 index 000000000000..853580a8a95e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_goal_selector.py @@ -0,0 +1,105 @@ +# 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 + +"""Minimal Newton-viewer selector for conveyor transfer goals.""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING, Any + +from .conveyor_geometry import CUBE_COLORS +from .mdp.reset_events import LEFT_SIDE + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def _mix_color( + color: tuple[float, float, float], + target: tuple[float, float, float], + amount: float, +) -> tuple[float, float, float]: + """Linearly mix two RGB colors.""" + return tuple(value + amount * (target_value - value) for value, target_value in zip(color, target, strict=True)) + + +class ConveyorGoalSelector: + """Render four color-matched cube buttons for one displayed environment.""" + + def __init__(self, env: ManagerBasedRLEnv, env_id: int) -> None: + """Create a selector bound to one vectorized environment index.""" + if not 0 <= env_id < env.num_envs: + raise IndexError(f"Conveyor selector environment {env_id} is out of range.") + self._env = env + self._env_id = env_id + self._command = env.command_manager.get_term("transfer") + self._target_cube_id = 0 + self._source_side_id = LEFT_SIDE + self._last_refresh_time = float("-inf") + + def render(self, imgui: Any) -> None: + """Draw the selector and publish a new transfer command when clicked.""" + imgui.set_next_item_open(True, imgui.Cond_.appearing) + if not imgui.collapsing_header(f"Transfer Goal · Env {self._env_id}"): + return + imgui.separator() + + if not self._refresh_command(): + imgui.text_disabled("Waiting for transfer state...") + return + + style = imgui.get_style() + available_width = float(imgui.get_content_region_avail().x) + spacing = float(style.item_spacing.x) + button_width = max(36.0, (available_width - spacing * (len(CUBE_COLORS) - 1)) / len(CUBE_COLORS)) + button_height = max(34.0, min(46.0, button_width * 0.72)) + + selected_cube_id: int | None = None + for cube_id, color in enumerate(CUBE_COLORS): + selected = cube_id == self._target_cube_id + text_color = (0.05, 0.05, 0.05) if sum(color) > 1.45 else (1.0, 1.0, 1.0) + border_color = (1.0, 0.88, 0.20) if selected else (0.12, 0.12, 0.12) + label = f"{cube_id + 1}##conveyor_goal_{cube_id}" + + imgui.push_style_color(imgui.Col_.button, imgui.ImVec4(*color, 1.0)) + imgui.push_style_color( + imgui.Col_.button_hovered, + imgui.ImVec4(*_mix_color(color, (1.0, 1.0, 1.0), 0.18), 1.0), + ) + imgui.push_style_color( + imgui.Col_.button_active, + imgui.ImVec4(*_mix_color(color, (0.0, 0.0, 0.0), 0.12), 1.0), + ) + imgui.push_style_color(imgui.Col_.border, imgui.ImVec4(*border_color, 1.0)) + imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(*text_color, 1.0)) + imgui.push_style_var(imgui.StyleVar_.frame_border_size, 3.0 if selected else 1.0) + imgui.push_style_var(imgui.StyleVar_.frame_rounding, 5.0) + + if imgui.button(label, imgui.ImVec2(button_width, button_height)): + selected_cube_id = cube_id + + imgui.pop_style_var(2) + imgui.pop_style_color(5) + if cube_id + 1 < len(CUBE_COLORS): + imgui.same_line() + + if selected_cube_id is not None and selected_cube_id != self._target_cube_id: + self._command.set_goal(selected_cube_id, env_ids=(self._env_id,)) + self._refresh_command(force=True) + + source_name = "Left" if self._source_side_id == LEFT_SIDE else "Right" + target_name = "Right" if self._source_side_id == LEFT_SIDE else "Left" + imgui.text(f"Cube {self._target_cube_id + 1}: {source_name} -> {target_name}") + imgui.text_disabled("Click a color to change the next transfer.") + + def _refresh_command(self, force: bool = False) -> bool: + """Refresh the small host-side UI cache at most ten times per second.""" + current_time = time.monotonic() + if force or current_time - self._last_refresh_time >= 0.1: + self._target_cube_id = int(self._command.target_cube_ids[self._env_id].item()) + self._source_side_id = int(self._command.source_side_ids[self._env_id].item()) + self._last_refresh_time = current_time + return True diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_warehouse_geometry.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_warehouse_geometry.py new file mode 100644 index 000000000000..17e42d8af016 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/conveyor_warehouse_geometry.py @@ -0,0 +1,226 @@ +# 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 + +"""Physical surfaces for the USD-authored elevated conveyor network.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from functools import lru_cache +from pathlib import Path + +import numpy as np + +from isaaclab.physics import SurfaceVelocitySpec + +from .conveyor_geometry import ( + BELT_CENTER_X, + BELT_CENTER_Y, + BELT_COLLISION_OVERHANG, + BELT_COLLISION_SEAM_OVERLAP, + BELT_HALF_STRAIGHT, + BELT_THICKNESS, + GUARD_BASE_OVERLAP, + GUARD_HEIGHT, + GUARD_THICKNESS, + ConveyorSectionSpec, + MeshSpec, + _turn_collision_mesh, + belt_collision_section_specs, + belt_direction, +) + +_ROUTE_ASSET = Path(__file__).parent / "assets" / "conveyor_routes.usda" + + +@lru_cache(maxsize=1) +def _route_layer(): + # Import Usd to register file formats in kitless mode, without composing asset dependencies. + from pxr import Sdf, Usd # noqa: F401 + + return Sdf.Layer.FindOrOpen(str(_ROUTE_ASSET)) + + +@dataclass(frozen=True) +class _RouteSegment: + name: str + points: tuple[tuple[float, float, float], ...] + direction: tuple[float, float, float] + width: float + pivot: tuple[float, float, float] | None + radius: float | None + friction: float | None + velocity: float | None + path: str + + +@lru_cache(maxsize=2) +def _route_segments(side: str) -> tuple[_RouteSegment, ...]: + """Read the same centerlines that author the visible conveyor modules.""" + from pxr import Sdf + + belt_direction(side) + layer = _route_layer() + root = layer.GetPrimAtPath(Sdf.Path(f"/ConveyorRoutes/{side}")) + segments = [] + for name in root.attributes["conveyor:order"].default: + prim = layer.GetPrimAtPath(root.path.AppendChild(name)) + points = layer.GetPrimAtPath(prim.path.AppendChild("Centerline")).attributes["points"].default + curved = "conveyor:pivot" in prim.attributes + segments.append( + _RouteSegment( + name=name, + velocity=prim.attributes["conveyor:velocity"].default + if "conveyor:velocity" in prim.attributes + else None, + path=prim.attributes["conveyor:path"].default if "conveyor:path" in prim.attributes else "main", + points=tuple(tuple(point) for point in points), + direction=tuple(prim.attributes["conveyor:direction"].default), + width=prim.attributes["conveyor:width"].default, + pivot=tuple(prim.attributes["conveyor:pivot"].default) if curved else None, + radius=prim.attributes["conveyor:radius"].default if curved else None, + friction=prim.attributes["conveyor:friction"].default + if "conveyor:friction" in prim.attributes + else None, + ) + ) + return tuple(segments) + + +def warehouse_parcel_positions() -> tuple[tuple[float, float, float], ...]: + """Return physical parcel infeed spawn positions in workspace coordinates [m].""" + root = _route_layer().GetPrimAtPath("/ConveyorRoutes") + return tuple(tuple(position) for position in root.attributes["conveyor:parcelSpawnPositions"].default) + + +def _normals(segment: _RouteSegment) -> np.ndarray: + if segment.pivot is None: + tangent = np.asarray(segment.direction) + normal = np.array((-tangent[1], tangent[0], 0.0)) + return np.tile(normal / np.linalg.norm(normal), (len(segment.points), 1)) + radial = np.asarray(segment.points) - segment.pivot + radial[:, 2] = 0 + return -math.copysign(1, segment.direction[2]) * radial / np.linalg.norm(radial, axis=1, keepdims=True) + + +def _prism( + name: str, + points: np.ndarray, + normals: np.ndarray, + offset: float | np.ndarray, + width: float, + bottom: float, + top: float, + *, + closed: bool = False, +) -> MeshSpec: + """Sweep a closed solid along a three-dimensional centerline [m].""" + offset = np.broadcast_to(offset, (len(points),))[:, None] + left = points + normals * (offset - width / 2) + right = points + normals * (offset + width / 2) + vertices = np.concatenate((left + (0, 0, top), right + (0, 0, top), left + (0, 0, bottom), right + (0, 0, bottom))) + count = len(points) + quads = [] + for i in range(count if closed else count - 1): + j = (i + 1) % count + quads.extend( + ( + (i, j, count + j, count + i), + (2 * count + i, 3 * count + i, 3 * count + j, 2 * count + j), + (i, 2 * count + i, 2 * count + j, j), + (count + i, count + j, 3 * count + j, 3 * count + i), + ) + ) + if not closed: + quads.extend(((0, count, 3 * count, 2 * count), (count - 1, 3 * count - 1, 4 * count - 1, 2 * count - 1))) + faces = tuple(triangle for a, b, c, d in quads for triangle in ((a, b, c), (a, c, d))) + return MeshSpec(name, tuple(tuple(vertex) for vertex in vertices), faces) + + +def warehouse_belt_sections(side: str, **kwargs: float | bool) -> tuple[ConveyorSectionSpec, ...]: + """Build linked belt surfaces from USD while preserving the original manipulation geometry.""" + sign = belt_direction(side) + sections = [] + for segment in _route_segments(side): + name = f"Conveyor{side}{segment.name}Collision" + pivot = segment.pivot + surface_kwargs = dict(kwargs) + if segment.velocity is not None: + surface_kwargs["velocity"] = segment.velocity + if segment.friction is not None: + surface_kwargs["friction_coefficient"] = segment.friction + if segment.name == "Working": + original = belt_collision_section_specs(side, **kwargs)[1 if sign > 0 else 0] + sections.append(original) + continue + if segment.name in {"PickupBend", "PlacementBend"}: + pickup = segment.name == "PickupBend" + name = f"Conveyor{side}{'Left' if pickup else 'Right'}InnerTurnCollision" + pivot = (BELT_CENTER_X + (-BELT_HALF_STRAIGHT if pickup else BELT_HALF_STRAIGHT), sign * BELT_CENTER_Y, 0.0) + geometry = _turn_collision_mesh( + name, pivot[0], BELT_CENTER_Y, math.pi if pickup else -math.pi / 2, math.pi / 2 + ) + if sign < 0: + geometry = replace( + geometry, + vertices=tuple((x, -y, z) for x, y, z in geometry.vertices), + faces=tuple((a, c, b) for a, b, c in geometry.faces), + ) + else: + points = np.array(segment.points) + # Overlap adjacent panels without exposing a vertical collision seam. + for index, neighbor, direction in ((0, 1, -1), (-1, -2, 1)): + tangent = points[neighbor] - points[index] if index == 0 else points[index] - points[neighbor] + points[index] += direction * BELT_COLLISION_SEAM_OVERLAP * tangent / np.linalg.norm(tangent) + geometry = _prism( + name, points, _normals(segment), 0, segment.width + 2 * BELT_COLLISION_OVERHANG, -BELT_THICKNESS, 0 + ) + sections.append( + ConveyorSectionSpec( + geometry, + SurfaceVelocitySpec( + prim_path=f"{{ENV_REGEX_NS}}/{geometry.name}", + direction=segment.direction, + surface_normal=(0.0, 0.0, 1.0) + if segment.pivot is not None + else tuple(np.cross(segment.direction, _normals(segment)[0])), + curved=segment.pivot is not None, + pivot_point=pivot or (0, 0, 0), + radius=segment.radius, + **surface_kwargs, + ), + ) + ) + return tuple(sections) + + +def warehouse_guard_meshes(side: str) -> tuple[MeshSpec, ...]: + """Build continuous guides for each closed circulation route and open supply belt.""" + segments = _route_segments(side) + guards = [] + for path in dict.fromkeys(segment.path for segment in segments): + points, normals, offsets = [], [], [] + route = [segment for segment in segments if segment.path == path] + for index, segment in enumerate(route): + stop = None if path != "main" and index == len(route) - 1 else -1 + for point, normal in zip(segment.points[:stop], _normals(segment)[:stop]): + points.append(point) + normals.append(normal) + offsets.append((segment.width + GUARD_THICKNESS) / 2) + for boundary, sign in (("Inner", -1), ("Outer", 1)): + guards.append( + _prism( + f"Guard{side}{path.title()}{boundary}", + np.asarray(points), + np.asarray(normals), + sign * np.asarray(offsets), + GUARD_THICKNESS, + -GUARD_BASE_OVERLAP, + GUARD_HEIGHT, + closed=path == "main", + ) + ) + return tuple(guards) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/franka_robot_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/franka_robot_cfg.py new file mode 100644 index 000000000000..12a0adba8338 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/franka_robot_cfg.py @@ -0,0 +1,69 @@ +# 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 + +"""Task-calibrated Franka configuration for conveyor manipulation.""" + +from isaaclab_newton.sim.schemas import MujocoJointCfg +from isaaclab_physx.sim.schemas import PhysxArticulationCfg + +from isaaclab.actuators import ImplicitActuatorCfg + +from isaaclab_assets.robots.franka import FRANKA_PANDA_MENAGERIE_CFG + +FRANKA_PANDA_CONVEYOR_CFG = FRANKA_PANDA_MENAGERIE_CFG.copy() +FRANKA_PANDA_CONVEYOR_CFG.spawn.rigid_props.disable_gravity = False +# Route gravity compensation through MuJoCo's actuator channel so effort limits +# and the solver apply it consistently with the configured implicit drives. +FRANKA_PANDA_CONVEYOR_CFG.spawn.joint_drive_props = [MujocoJointCfg(actuatorgravcomp=True)] +FRANKA_PANDA_CONVEYOR_CFG.actuators = { + "panda_arm": ImplicitActuatorCfg( + joint_names_expr=["panda_joint[1-7]"], + effort_limit_sim={"panda_joint[1-4]": 87.0, "panda_joint[5-7]": 12.0}, + velocity_limit_sim={"panda_joint[1-4]": 20.0, "panda_joint[5-7]": 25.0}, + stiffness={ + "panda_joint[1-4]": 600.0, + "panda_joint5": 250.0, + "panda_joint6": 150.0, + "panda_joint7": 50.0, + }, + damping={ + "panda_joint[1-4]": 50.0, + "panda_joint5": 30.0, + "panda_joint6": 25.0, + "panda_joint7": 15.0, + }, + armature={ + "panda_joint[1-2]": 0.6057, + "panda_joint[3-4]": 0.4625, + "panda_joint[5-7]": 0.2055, + }, + ), + "panda_hand": ImplicitActuatorCfg( + joint_names_expr=["panda_finger_joint[1-2]"], + effort_limit_sim=70.0, + velocity_limit_sim=2.0, + stiffness=350.0, + # Keep the 0.1 kg m^2 armature close to critical damping so the + # fingers establish contact within a few 50 Hz policy steps. + damping=10.0, + armature=0.1, + ), +} +"""Menagerie Franka with explicit manipulation gains and solver-native gravity compensation.""" + + +FRANKA_PANDA_CONVEYOR_PHYSX_CFG = FRANKA_PANDA_CONVEYOR_CFG.copy() +# PhysX does not consume MuJoCo's actuator-gravity-compensation attribute. Disabling +# gravity on the robot is the closest solver-native equivalent and keeps the trained +# position-policy contract unchanged without adding a task-side effort loop. +FRANKA_PANDA_CONVEYOR_PHYSX_CFG.spawn.rigid_props.disable_gravity = True +FRANKA_PANDA_CONVEYOR_PHYSX_CFG.spawn.joint_drive_props = None +# Contact-rich manipulation benefits from resolving the articulation for more than +# the generic asset defaults, especially with the deliberately stiff trained gains. +for _properties in FRANKA_PANDA_CONVEYOR_PHYSX_CFG.spawn.articulation_props: + if isinstance(_properties, PhysxArticulationCfg): + _properties.solver_position_iteration_count = 32 + _properties.solver_velocity_iteration_count = 4 +"""PhysX variant with the same joints, gains, action ordering, and gravity-compensated policy contract.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.py new file mode 100644 index 000000000000..56b9c4494587 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""MDP terms for the conveyor-to-conveyor Franka transfer task.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.pyi new file mode 100644 index 000000000000..5bd90e830c86 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/__init__.pyi @@ -0,0 +1,79 @@ +# 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 + +__all__ = [ + "BELT_DEPLOYMENT_VARIANT", + "ConveyorRelativeJointPositionAction", + "ConveyorRelativeJointPositionActionCfg", + "ConveyorResetCurriculum", + "ConveyorResetRecipe", + "ConveyorResetStateTable", + "ConveyorTransferCommand", + "ConveyorTransferCommandCfg", + "ResetBufferedGripperAction", + "ResetBufferedGripperActionCfg", + "SuccessMonitorCfg", + "action_term_l2", + "active_transfer_features", + "build_reset_rows", + "cube_conveyor_state", + "cube_out_of_workspace", + "end_effector_axes", + "end_effector_velocity", + "finite_action_rate_l2", + "finite_joint_velocity_l2", + "gripper_joint_positions", + "invalid_action", + "nonfinite_scene_state", + "physical_cube_acquisition_mask", + "select_next_transfer_cube", + "subgoal_time_out", + "target_cube_one_hot", + "target_side_one_hot", + "terminal_failure", + "transfer_object_observation", + "transfer_sequence_time_out", + "transfer_success_mask", + "transfer_success_reward", +] + +from .actions import ConveyorRelativeJointPositionAction, ResetBufferedGripperAction +from .actions_cfg import ConveyorRelativeJointPositionActionCfg, ResetBufferedGripperActionCfg +from .commands import ConveyorTransferCommand, ConveyorTransferCommandCfg, transfer_success_mask +from .curriculums import ConveyorResetCurriculum +from .observations import ( + active_transfer_features, + cube_conveyor_state, + end_effector_axes, + end_effector_velocity, + gripper_joint_positions, + target_cube_one_hot, + target_side_one_hot, + transfer_object_observation, +) +from .reset_events import ( + BELT_DEPLOYMENT_VARIANT, + ConveyorResetRecipe, + ConveyorResetStateTable, + build_reset_rows, + select_next_transfer_cube, +) +from .rewards import ( + action_term_l2, + finite_action_rate_l2, + finite_joint_velocity_l2, + physical_cube_acquisition_mask, + terminal_failure, + transfer_success_reward, +) +from .terminations import ( + cube_out_of_workspace, + invalid_action, + nonfinite_scene_state, + subgoal_time_out, + transfer_sequence_time_out, +) +from isaaclab.envs.mdp import * +from isaaclab_tasks.utils.success_monitor import SuccessMonitorCfg diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions.py new file mode 100644 index 000000000000..bdadbac2b7b7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions.py @@ -0,0 +1,135 @@ +# 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 + +"""Reset-safe relative joint actions for conveyor transfer.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.envs.mdp.actions.binary_joint_actions import BinaryJointPositionAction +from isaaclab.envs.mdp.actions.joint_actions import JointAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .actions_cfg import ConveyorRelativeJointPositionActionCfg, ResetBufferedGripperActionCfg + + +class ConveyorRelativeJointPositionAction(JointAction): + """Apply one measured-state-relative target per policy step. + + Isaac Lab's generic relative action adds the residual during every physics + substep. This term computes the target once in :meth:`process_actions`, so + simulation decimation does not multiply the requested displacement. + """ + + cfg: ConveyorRelativeJointPositionActionCfg + + def __init__(self, cfg: ConveyorRelativeJointPositionActionCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + if not math.isfinite(cfg.max_delta) or cfg.max_delta <= 0.0: + raise ValueError("max_delta must be finite and positive.") + if not math.isfinite(cfg.joint_limit_margin) or cfg.joint_limit_margin < 0.0: + raise ValueError("joint_limit_margin must be finite and non-negative.") + self._workspace_lower = torch.tensor(cfg.workspace_lower, dtype=torch.float32, device=self.device) + self._workspace_upper = torch.tensor(cfg.workspace_upper, dtype=torch.float32, device=self.device) + if self._workspace_lower.shape != (self.action_dim,) or self._workspace_upper.shape != (self.action_dim,): + raise ValueError("workspace bounds must contain one value per controlled joint.") + if torch.any(self._workspace_lower >= self._workspace_upper): + raise ValueError("Every lower workspace bound must be less than its upper bound.") + if not torch.all(torch.isfinite(self._workspace_lower)) or not torch.all(torch.isfinite(self._workspace_upper)): + raise ValueError("workspace bounds must be finite.") + self._previous_actions = torch.zeros_like(self._raw_actions) + self._invalid_actions = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self._position_targets = self._asset.data.joint_pos.torch[:, self._joint_ids].clone() + + @property + def previous_actions(self) -> torch.Tensor: + """Previous finite policy actions, shape ``(num_envs, action_dim)``.""" + return self._previous_actions + + @property + def invalid_actions(self) -> torch.Tensor: + """Whether the latest policy action contained a non-finite component.""" + return self._invalid_actions + + def process_actions(self, actions: torch.Tensor) -> None: + """Convert normalized residuals into bounded position targets [rad].""" + self._previous_actions.copy_(self._raw_actions) + self._invalid_actions.copy_(~torch.isfinite(actions).all(dim=1)) + finite_actions = torch.nan_to_num(actions, nan=0.0, posinf=1.0, neginf=-1.0).clamp(-1.0, 1.0) + super().process_actions(finite_actions) + delta = torch.clamp(self._processed_actions, min=-self.cfg.max_delta, max=self.cfg.max_delta) + positions = self._asset.data.joint_pos.torch[:, self._joint_ids] + limits = self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids] + lower = torch.maximum(limits[..., 0] + self.cfg.joint_limit_margin, self._workspace_lower) + upper = torch.minimum(limits[..., 1] - self.cfg.joint_limit_margin, self._workspace_upper) + self._position_targets = torch.clamp(positions + delta, min=lower, max=upper) + self._processed_actions = self._position_targets + + def apply_actions(self) -> None: + """Hold the policy-step target through all physics substeps.""" + self._asset.set_joint_position_target_index(target=self._position_targets, joint_ids=self._joint_ids) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Initialize targets from the sampled reset pose.""" + super().reset(env_ids) + positions = self._asset.data.joint_pos.torch[:, self._joint_ids] + if env_ids is None: + self._position_targets[:] = positions + self._processed_actions[:] = positions + else: + self._position_targets[env_ids] = positions[env_ids] + self._processed_actions[env_ids] = positions[env_ids] + self._previous_actions[env_ids] = 0.0 + self._invalid_actions[env_ids] = False + + +class ResetBufferedGripperAction(BinaryJointPositionAction): + """Keep reset-authored grasps closed during a short settling window.""" + + cfg: ResetBufferedGripperActionCfg + + def __init__(self, cfg: ResetBufferedGripperActionCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + self._previous_actions = torch.zeros_like(self._raw_actions) + self._invalid_actions = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + @property + def previous_actions(self) -> torch.Tensor: + """Previous finite policy actions, shape ``(num_envs, action_dim)``.""" + return self._previous_actions + + @property + def invalid_actions(self) -> torch.Tensor: + """Whether the latest gripper command contained a non-finite value.""" + return self._invalid_actions + + def process_actions(self, actions: torch.Tensor) -> None: + """Map finite binary commands and preserve initially held cubes.""" + self._previous_actions.copy_(self._raw_actions) + if actions.dtype == torch.bool: + self._invalid_actions.zero_() + finite_actions = actions + else: + self._invalid_actions.copy_(~torch.isfinite(actions).all(dim=1)) + # A non-finite gripper command closes the fingers for the final safe + # step before the invalid-action termination is evaluated. + finite_actions = torch.nan_to_num(actions, nan=-1.0, posinf=1.0, neginf=-1.0).clamp(-1.0, 1.0) + super().process_actions(finite_actions) + command = self._env.command_manager.get_term(self.cfg.command_name) + force_close = (command.held_cube_ids >= 0) & (self._env.episode_length_buf < self.cfg.force_close_steps) + self._processed_actions[force_close] = self._close_command + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Clear buffered invalid-command state for selected environments.""" + super().reset(env_ids) + self._previous_actions[env_ids] = 0.0 + self._invalid_actions[env_ids] = False diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions_cfg.py new file mode 100644 index 000000000000..40bc923f779e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/actions_cfg.py @@ -0,0 +1,48 @@ +# 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 + +"""Action configurations for conveyor transfer.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.envs.mdp.actions.actions_cfg import BinaryJointPositionActionCfg, JointActionCfg +from isaaclab.utils.configclass import configclass + +if TYPE_CHECKING: + from .actions import ConveyorRelativeJointPositionAction, ResetBufferedGripperAction + + +@configclass +class ConveyorRelativeJointPositionActionCfg(JointActionCfg): + """Configuration for measured-state relative Franka joint control.""" + + class_type: type[ConveyorRelativeJointPositionAction] | str = "{DIR}.actions:ConveyorRelativeJointPositionAction" + + joint_limit_margin: float = 0.02 + """Distance kept from each soft joint limit [rad].""" + + max_delta: float = 0.12 + """Maximum target change per policy step [rad].""" + + workspace_lower: tuple[float, ...] = (-0.75, -0.45, -0.55, -2.75, -0.45, 1.85, -0.10) + """Lower boundary of the validated transfer workspace [rad].""" + + workspace_upper: tuple[float, ...] = (0.85, 0.85, 0.35, -1.75, 0.45, 3.05, 1.65) + """Upper boundary of the validated transfer workspace [rad].""" + + +@configclass +class ResetBufferedGripperActionCfg(BinaryJointPositionActionCfg): + """Configuration for reset-grasp protection.""" + + class_type: type[ResetBufferedGripperAction] | str = "{DIR}.actions:ResetBufferedGripperAction" + + force_close_steps: int = 5 + """Initial policy steps that preserve a reset-authored grasp.""" + + command_name: str = "transfer" + """Command term that owns reset-authored held-cube state.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/commands.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/commands.py new file mode 100644 index 000000000000..2c03810c417f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/commands.py @@ -0,0 +1,348 @@ +# 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 + +"""Success-driven transfer commands for the conveyor Franka task.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import CommandTerm, CommandTermCfg +from isaaclab.utils.configclass import configclass + +from ..conveyor_cube_pool import cube_values +from ..conveyor_geometry import BELT_CENTER_X, BELT_HALF_STRAIGHT +from .kinematics import end_effector_pose +from .reset_events import CUBE_COUNT, ConveyorResetRecipe, select_next_transfer_cube, side_inner_y +from .rewards import current_transfer_potential, physical_cube_acquisition_mask + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +def transfer_success_mask( + cube_positions: torch.Tensor, + cube_linear_velocities: torch.Tensor, + tool_positions: torch.Tensor, + finger_positions: torch.Tensor, + target_side_ids: torch.Tensor, + lateral_tolerance: float = 0.055, + maximum_cube_speed: float = 0.65, + minimum_finger_position: float = 0.027, + minimum_tool_clearance: float = 0.055, +) -> torch.Tensor: + """Return whether the active cube is released on its destination belt.""" + target_y = side_inner_y(target_side_ids) + on_straight = torch.abs(cube_positions[:, 0] - BELT_CENTER_X) < BELT_HALF_STRAIGHT + on_lane = torch.abs(cube_positions[:, 1] - target_y) < lateral_tolerance + supported_height = (cube_positions[:, 2] > 0.045) & (cube_positions[:, 2] < 0.095) + moving_safely = torch.linalg.vector_norm(cube_linear_velocities, dim=1) < maximum_cube_speed + released = torch.amin(finger_positions, dim=1) > minimum_finger_position + hand_clear = torch.linalg.vector_norm(tool_positions - cube_positions, dim=1) > minimum_tool_clearance + return on_straight & on_lane & supported_height & moving_safely & released & hand_clear + + +class ConveyorTransferCommand(CommandTerm): + """Command one numbered cube to the opposite belt and redraw on success.""" + + cfg: ConveyorTransferCommandCfg + + def __init__(self, cfg: ConveyorTransferCommandCfg, env: ManagerBasedRLEnv) -> None: + self._validate_cfg(cfg) + super().__init__(cfg, env) + + reset_term = env.event_manager.get_term_cfg(cfg.reset_event_name).func + required_reset_fields = ("row_ids", "recipe_ids", "target_cube_ids", "source_side_ids", "held_rows") + if not all(hasattr(reset_term, name) for name in required_reset_fields): + raise RuntimeError("ConveyorTransferCommand requires ConveyorResetStateTable reset metadata.") + self._reset_term = reset_term + self._robot: Articulation = env.scene["robot"] + self._finger_joint_ids = self._robot.find_joints("panda_finger_joint[1-2]", preserve_order=True)[0] + + self.target_cube_ids = torch.zeros(self.num_envs, dtype=torch.long, device=self.device) + self.source_side_ids = torch.zeros_like(self.target_cube_ids) + self.recipe_ids = torch.zeros_like(self.target_cube_ids) + self.held_cube_ids = torch.full_like(self.target_cube_ids, -1) + self.subgoal_start_steps = torch.zeros_like(self.target_cube_ids) + self.transfer_counts = torch.zeros_like(self.target_cube_ids) + self.direction_transfer_counts = torch.zeros((self.num_envs, 2), dtype=torch.long, device=self.device) + + self._stable_steps = torch.zeros_like(self.target_cube_ids) + self.is_success = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self.new_success = torch.zeros_like(self.is_success) + self.pending_success = torch.zeros_like(self.is_success) + self.ever_success = torch.zeros_like(self.is_success) + self.progress_ever_success = torch.zeros_like(self.is_success) + self._target_potential = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device) + self._last_evaluation_steps = torch.full_like(self.target_cube_ids, -1) + self._resampling_from_reset = False + + self.metrics["success_rate"] = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device) + self.metrics["transfer_count"] = torch.zeros_like(self.metrics["success_rate"]) + self.metrics["left_to_right_transfers"] = torch.zeros_like(self.metrics["success_rate"]) + self.metrics["right_to_left_transfers"] = torch.zeros_like(self.metrics["success_rate"]) + + @staticmethod + def _validate_cfg(cfg: ConveyorTransferCommandCfg) -> None: + """Validate success and reset-progress thresholds.""" + if cfg.minimum_subgoal_steps < 0 or cfg.hold_steps < 1: + raise ValueError("minimum_subgoal_steps must be non-negative and hold_steps must be positive.") + if cfg.lateral_tolerance <= 0.0 or cfg.maximum_cube_speed <= 0.0: + raise ValueError("Conveyor placement tolerances and speed limits must be positive.") + if cfg.minimum_finger_position <= 0.0 or cfg.minimum_tool_clearance <= 0.0: + raise ValueError("Conveyor release thresholds must be positive.") + if cfg.minimum_progress_steps < 0 or cfg.minimum_progress <= 0.0 or cfg.maximum_target_potential <= 0.0: + raise ValueError("Conveyor reset-progress thresholds are invalid.") + if ( + cfg.minimum_acquisition_lift <= 0.0 + or cfg.maximum_acquisition_tool_distance <= 0.0 + or cfg.maximum_acquisition_finger_position <= 0.0 + ): + raise ValueError("Conveyor acquisition thresholds must be positive.") + + @property + def command(self) -> torch.Tensor: + """Return target-cube and destination-belt one-hot commands.""" + cube = torch.nn.functional.one_hot(self.target_cube_ids, num_classes=CUBE_COUNT) + destination = torch.nn.functional.one_hot(1 - self.source_side_ids, num_classes=2) + return torch.cat((cube, destination), dim=1).float() + + def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, float]: + """Log per-command completion and initialize commands from sampled reset rows.""" + ids = self._resolve_env_ids(env_ids) + valid = self.command_counter[ids] > 0 + attempts = self.command_counter[ids].clamp_min(1) + self.metrics["success_rate"][ids] = torch.where( + valid, + self.transfer_counts[ids].float() / attempts.float(), + 0.0, + ) + self.metrics["transfer_count"][ids] = self.transfer_counts[ids].float() + self.metrics["left_to_right_transfers"][ids] = self.direction_transfer_counts[ids, 0].float() + self.metrics["right_to_left_transfers"][ids] = self.direction_transfer_counts[ids, 1].float() + + self._resampling_from_reset = True + try: + extras = super().reset(ids) + finally: + self._resampling_from_reset = False + + initial_potential = current_transfer_potential(self._env, command=self) + self._target_potential[ids] = torch.clamp_max( + initial_potential[ids] + self.cfg.minimum_progress, + self.cfg.maximum_target_potential, + ) + self._last_evaluation_steps[ids] = -1 + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = extras.pop("success_rate") + return extras + + def evaluate(self) -> None: + """Evaluate stable completion and reset-learning progress once per policy step.""" + evaluation_steps = self._env.episode_length_buf + evaluate_mask = (self.command_counter > 0) & (self._last_evaluation_steps != evaluation_steps) + if not bool(torch.any(evaluate_mask)): + return + + positions = cube_values(self._env, "root_pos_w") + velocities = cube_values(self._env, "root_lin_vel_w") + index = self.target_cube_ids.view(self.num_envs, 1, 1).expand(-1, 1, 3) + active_position = torch.gather(positions, 1, index).squeeze(1) - self._env.scene.env_origins + active_velocity = torch.gather(velocities, 1, index).squeeze(1) + tool_position, _ = end_effector_pose(self._env) + tool_position = tool_position - self._env.scene.env_origins + finger_positions = self._robot.data.joint_pos.torch[:, self._finger_joint_ids] + successful = transfer_success_mask( + active_position, + active_velocity, + tool_position, + finger_positions, + 1 - self.source_side_ids, + lateral_tolerance=self.cfg.lateral_tolerance, + maximum_cube_speed=self.cfg.maximum_cube_speed, + minimum_finger_position=self.cfg.minimum_finger_position, + minimum_tool_clearance=self.cfg.minimum_tool_clearance, + ) + subgoal_steps = evaluation_steps - self.subgoal_start_steps + successful &= subgoal_steps >= self.cfg.minimum_subgoal_steps + + next_stable_steps = torch.where(successful, self._stable_steps + 1, torch.zeros_like(self._stable_steps)) + self._stable_steps[evaluate_mask] = next_stable_steps[evaluate_mask] + stable = self._stable_steps >= self.cfg.hold_steps + new_success = evaluate_mask & stable & ~self.is_success & ~self.pending_success + self.is_success[evaluate_mask] = stable[evaluate_mask] + self.new_success.copy_(new_success) + self.pending_success |= new_success + self.ever_success |= new_success + success_ids = new_success.nonzero(as_tuple=False).squeeze(-1) + if success_ids.numel(): + source_sides = self.source_side_ids[success_ids] + self.transfer_counts[success_ids] += 1 + self.direction_transfer_counts[success_ids, source_sides] += 1 + pool = getattr(self._env, "conveyor_cube_pool", None) + if pool is not None: + pool.record_transfers(success_ids, self.target_cube_ids[success_ids]) + + potential = current_transfer_potential(self._env, command=self) + progressed = (potential >= self._target_potential) & (evaluation_steps >= self.cfg.minimum_progress_steps) + acquisition_recipe = ( + (self.recipe_ids == int(ConveyorResetRecipe.GRASP)) + | (self.recipe_ids == int(ConveyorResetRecipe.PREGRASP)) + | (self.recipe_ids == int(ConveyorResetRecipe.BELT)) + ) + physically_acquired = physical_cube_acquisition_mask( + self._env, + command=self, + minimum_lift=self.cfg.minimum_acquisition_lift, + maximum_tool_distance=self.cfg.maximum_acquisition_tool_distance, + maximum_finger_position=self.cfg.maximum_acquisition_finger_position, + ) + progressed &= ~acquisition_recipe | physically_acquired + self.progress_ever_success |= evaluate_mask & progressed + self._last_evaluation_steps[evaluate_mask] = evaluation_steps[evaluate_mask] + self._env.extras["successes"] = self.ever_success.clone() + + def set_goal(self, target_cube_id: int, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: + """Replace the active command using the selected cube's current conveyor.""" + if not isinstance(target_cube_id, int) or isinstance(target_cube_id, bool): + raise TypeError("target_cube_id must be an integer.") + if not 0 <= target_cube_id < CUBE_COUNT: + raise ValueError(f"target_cube_id must lie in [0, {CUBE_COUNT - 1}].") + ids = self._resolve_env_ids(env_ids) + if ids.numel() == 0: + return + local_y = cube_values(self._env, "root_pos_w")[ids, target_cube_id, 1] - self._env.scene.env_origins[ids, 1] + source_side_ids = (local_y < 0.0).long() + target_cube_ids = torch.full_like(ids, target_cube_id) + self._assign_goal(ids, target_cube_ids, source_side_ids) + + def _update_metrics(self) -> None: + self.evaluate() + + def _resample_command(self, env_ids: Sequence[int]) -> None: + ids = self._resolve_env_ids(env_ids) + if ids.numel() == 0: + return + if self._resampling_from_reset: + rows = self._reset_term.row_ids[ids] + target_cube_ids = self._reset_term.target_cube_ids[rows] + source_side_ids = self._reset_term.source_side_ids[rows] + self.recipe_ids[ids] = self._reset_term.recipe_ids[rows] + held_rows = self._reset_term.held_rows[rows] + self.transfer_counts[ids] = 0 + self.direction_transfer_counts[ids] = 0 + self.ever_success[ids] = False + self.progress_ever_success[ids] = False + self._assign_goal(ids, target_cube_ids, source_side_ids) + self.held_cube_ids[ids] = torch.where(held_rows, target_cube_ids, -1) + self.subgoal_start_steps[ids] = 0 + return + + positions = cube_values(self._env, "root_pos_w")[ids] + positions -= self._env.scene.env_origins[ids].unsqueeze(1) + next_source_side_ids = 1 - self.source_side_ids[ids] + next_cube_ids = select_next_transfer_cube( + positions, + self.target_cube_ids[ids], + next_source_side_ids, + transit_half_width=self.cfg.transit_half_width, + ) + self._assign_goal(ids, next_cube_ids, next_source_side_ids) + + def _update_command(self) -> None: + completed_ids = self.pending_success.nonzero(as_tuple=False).squeeze(-1) + if completed_ids.numel(): + self._resample(completed_ids) + + def _assign_goal( + self, + env_ids: torch.Tensor, + target_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, + ) -> None: + """Publish one command and clear completion state from the previous command.""" + self.target_cube_ids[env_ids] = target_cube_ids + self.source_side_ids[env_ids] = source_side_ids + self.held_cube_ids[env_ids] = -1 + self.subgoal_start_steps[env_ids] = self._env.episode_length_buf[env_ids] + self._stable_steps[env_ids] = 0 + self.is_success[env_ids] = False + self.new_success[env_ids] = False + self.pending_success[env_ids] = False + self._last_evaluation_steps[env_ids] = -1 + + def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | slice | None) -> torch.Tensor: + """Return validated environment indices on the command device.""" + if env_ids is None: + ids = torch.arange(self.num_envs, dtype=torch.long, device=self.device) + elif isinstance(env_ids, slice): + ids = torch.arange(self.num_envs, dtype=torch.long, device=self.device)[env_ids] + else: + ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.device).flatten() + if bool(torch.any((ids < 0) | (ids >= self.num_envs))): + raise IndexError("Conveyor command environment indices are out of range.") + return ids + + def _set_debug_vis_impl(self, debug_vis: bool) -> None: + raise NotImplementedError("Conveyor commands are visualized by the Newton goal selector.") + + def _debug_vis_callback(self, event) -> None: + raise NotImplementedError("Conveyor commands are visualized by the Newton goal selector.") + + +@configclass +class ConveyorTransferCommandCfg(CommandTermCfg): + """Configuration for success-driven, bidirectional cube-transfer commands.""" + + class_type: type[ConveyorTransferCommand] | str = "{DIR}.commands:ConveyorTransferCommand" + """Command-term implementation resolved lazily by the manager.""" + + resampling_time_range: tuple[float, float] = (1.0e6, 1.0e6) + """Time-based resampling interval [s]; success normally resamples first.""" + + reset_event_name: str = "reset_from_state_table" + """Reset event containing the immutable physical-state table.""" + + transit_half_width: float = 0.14 + """Half-width of the corridor excluded when selecting a source-belt cube [m].""" + + minimum_subgoal_steps: int = 2 + """Minimum policy steps before a placement can succeed.""" + + hold_steps: int = 3 + """Consecutive policy steps for which a placement must remain valid.""" + + lateral_tolerance: float = 0.055 + """Maximum lateral error from the destination belt center [m].""" + + maximum_cube_speed: float = 0.65 + """Maximum cube speed for a completed placement [m/s].""" + + minimum_finger_position: float = 0.027 + """Minimum position of each finger for the cube to count as released [m].""" + + minimum_tool_clearance: float = 0.055 + """Minimum tool-to-cube distance for the cube to count as released [m].""" + + minimum_progress_steps: int = 3 + """Minimum policy steps before reset-learning progress can be credited.""" + + minimum_progress: float = 0.35 + """Required increase in the dimensionless transfer potential after reset.""" + + maximum_target_potential: float = 5.0 + """Upper bound for the dimensionless reset-progress target.""" + + minimum_acquisition_lift: float = 0.025 + """Minimum cube lift used to validate physical acquisition [m].""" + + maximum_acquisition_tool_distance: float = 0.075 + """Maximum tool-to-cube distance used to validate physical acquisition [m].""" + + maximum_acquisition_finger_position: float = 0.030 + """Maximum finger position used to validate physical acquisition [m].""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/curriculums.py new file mode 100644 index 000000000000..85b0830ea909 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/curriculums.py @@ -0,0 +1,301 @@ +# 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 + +"""Adaptive reset-state curriculum for conveyor transfer.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import CurriculumTermCfg, ManagerTermBase + +from isaaclab_tasks.utils.success_monitor import SuccessMonitorCfg + +from .reset_events import BELT_DEPLOYMENT_VARIANT, CUBE_COUNT, ConveyorResetRecipe, reset_variant_counts + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def reset_sampling_probabilities( + recipe_ids: torch.Tensor, + variant_ids: torch.Tensor, + target_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, + target_weights: torch.Tensor, + deployment_probability: float | torch.Tensor, +) -> torch.Tensor: + """Balance target-rate weights and mix in guaranteed deployment starts.""" + if not ( + recipe_ids.shape == variant_ids.shape == target_cube_ids.shape == source_side_ids.shape == target_weights.shape + ): + raise ValueError("Reset row metadata and target weights must have matching shapes.") + deployment_probability = torch.as_tensor( + deployment_probability, + dtype=torch.float32, + device=target_weights.device, + ) + if deployment_probability.numel() != 1: + raise ValueError("deployment_probability must be a scalar.") + deployment_probability = deployment_probability.reshape(()) + if bool((deployment_probability <= 0.0) | (deployment_probability >= 1.0)): + raise ValueError("deployment_probability must lie strictly between zero and one.") + if not bool(torch.all(torch.isfinite(target_weights))) or bool(torch.any(target_weights < 0.0)): + raise ValueError("target_weights must be finite and non-negative.") + + deployment_rows = (recipe_ids == int(ConveyorResetRecipe.BELT)) & (variant_ids == BELT_DEPLOYMENT_VARIANT) + if not bool(torch.any(deployment_rows)) or bool(torch.all(deployment_rows)): + raise ValueError("Reset table must contain deployment and intermediate rows.") + + adaptive = target_weights.clone() + adaptive[deployment_rows] = 0.0 + + # Mirror Franka Stack's recipe/layout balancing: success in one physical + # phase must not starve another cube identity or transfer direction. + command_ids = 2 * target_cube_ids + source_side_ids + command_count = 2 * CUBE_COUNT + if bool(torch.any((command_ids < 0) | (command_ids >= command_count))): + raise ValueError("Reset table contains an invalid cube or source-side id.") + stratum_ids = recipe_ids * command_count + command_ids + stratum_count = len(ConveyorResetRecipe) * command_count + stratum_mass = torch.zeros(stratum_count, dtype=adaptive.dtype, device=adaptive.device) + stratum_mass.scatter_add_(0, stratum_ids, adaptive) + if bool(torch.any(stratum_mass <= 0.0)): + raise ValueError("Every recipe, cube, and source-side stratum must have intermediate reset rows.") + adaptive /= stratum_mass[stratum_ids] + adaptive[deployment_rows] = 0.0 + adaptive *= (1.0 - deployment_probability) / adaptive.sum() + deployment = deployment_rows.to(dtype=adaptive.dtype) + deployment *= deployment_probability / deployment.sum() + return adaptive + deployment + + +def deployment_probability_from_progress( + progress_rate: torch.Tensor, + row_coverage: torch.Tensor, + initial_probability: float = 0.35, + final_probability: float = 0.90, + progress_start: float = 0.45, + progress_end: float = 0.80, + coverage_target: float = 0.50, +) -> torch.Tensor: + """Interpolate deployment sampling from rolling competence and row coverage.""" + if progress_rate.numel() != 1 or row_coverage.numel() != 1: + raise ValueError("progress_rate and row_coverage must be scalar tensors.") + if not 0.0 < initial_probability <= final_probability < 1.0: + raise ValueError("Deployment probabilities must be ordered strictly inside (0, 1).") + if not 0.0 <= progress_start < progress_end <= 1.0: + raise ValueError("Deployment progress thresholds must be ordered inside [0, 1].") + if not 0.0 < coverage_target <= 1.0: + raise ValueError("coverage_target must lie inside (0, 1].") + if bool((progress_rate < 0.0) | (progress_rate > 1.0) | (row_coverage < 0.0) | (row_coverage > 1.0)): + raise ValueError("Rolling progress and row coverage must lie inside [0, 1].") + + progress_fraction = ((progress_rate - progress_start) / (progress_end - progress_start)).clamp(0.0, 1.0) + coverage_fraction = (row_coverage / coverage_target).clamp(0.0, 1.0) + readiness = progress_fraction * coverage_fraction + readiness = readiness.square() * (3.0 - 2.0 * readiness) + return initial_probability + (final_probability - initial_probability) * readiness + + +class ConveyorResetCurriculum(ManagerTermBase): + """Record row outcomes and sample the next physical reset states.""" + + def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + reset_term = env.event_manager.get_term_cfg("reset_from_state_table").func + if not hasattr(reset_term, "row_count"): + raise RuntimeError("ConveyorResetCurriculum requires ConveyorResetStateTable.") + self._reset_term = reset_term + self._attempts = torch.zeros(reset_term.row_count, dtype=torch.long, device=env.device) + self._progress_successes = torch.zeros_like(self._attempts) + self._final_successes = torch.zeros_like(self._attempts) + monitor_cfg = cfg.params.get("success_monitor") + if not isinstance(monitor_cfg, SuccessMonitorCfg): + raise TypeError("ConveyorResetCurriculum requires a SuccessMonitorCfg.") + self._progress_monitor = monitor_cfg.class_type( + monitor_cfg, + num_partitions=1, + partition_size=reset_term.row_count, + device=env.device, + ) + variant_counts = reset_variant_counts() + self._diagnostic_variant_rows = tuple( + ( + recipe.name.lower(), + variant_id, + (reset_term.recipe_ids == int(recipe)) & (reset_term.variant_ids == variant_id), + ) + for recipe in (ConveyorResetRecipe.PREGRASP, ConveyorResetRecipe.BELT) + for variant_id in range(variant_counts[int(recipe)]) + ) + + def __call__( + self, + env: ManagerBasedRLEnv, + env_ids: Sequence[int], + success_monitor: SuccessMonitorCfg, + command_name: str = "transfer", + deployment_probability_initial: float = 0.35, + deployment_probability_final: float = 0.90, + deployment_progress_start: float = 0.45, + deployment_progress_end: float = 0.80, + deployment_coverage_target: float = 0.50, + fixed_source_side_id: int | None = None, + ) -> dict[str, torch.Tensor]: + """Update adaptive evidence, sample rows, and expose diagnostics.""" + del success_monitor + ids = torch.as_tensor(env_ids, dtype=torch.long, device=env.device).flatten() + command = env.command_manager.get_term(command_name) + batch_progress = torch.zeros((), dtype=torch.float32, device=env.device) + batch_success = torch.zeros((), dtype=torch.float32, device=env.device) + completed = (command.command_counter[ids] > 0) & (env.episode_length_buf[ids] > 0) + completed_ids = ids[completed] + if completed_ids.numel(): + progressed = command.progress_ever_success[completed_ids] + succeeded = command.ever_success[completed_ids] + rows = self._reset_term.row_ids[completed_ids] + self._progress_monitor.success_update(rows, progressed) + self._attempts.scatter_add_(0, rows, torch.ones_like(rows)) + self._progress_successes.scatter_add_(0, rows, progressed.long()) + self._final_successes.scatter_add_(0, rows, succeeded.long()) + batch_progress = progressed.float().mean() + batch_success = succeeded.float().mean() + + history_size = self._progress_monitor.success_size + history_success_count = self._progress_monitor.success_buf.sum(dim=1) + attempted_rows = history_size > 0 + row_coverage = attempted_rows.float().mean() + total_progress = history_success_count.sum() / history_size.sum().clamp_min(1) + deployment_probability = deployment_probability_from_progress( + total_progress, + row_coverage, + initial_probability=deployment_probability_initial, + final_probability=deployment_probability_final, + progress_start=deployment_progress_start, + progress_end=deployment_progress_end, + coverage_target=deployment_coverage_target, + ) + probabilities = reset_sampling_probabilities( + self._reset_term.recipe_ids, + self._reset_term.variant_ids, + self._reset_term.target_cube_ids, + self._reset_term.source_side_ids, + self._progress_monitor.target_weights(), + deployment_probability, + ) + if fixed_source_side_id is not None: + if fixed_source_side_id not in (0, 1): + raise ValueError("fixed_source_side_id must be 0 (left) or 1 (right).") + probabilities *= self._reset_term.source_side_ids == fixed_source_side_id + probabilities /= probabilities.sum() + if ids.numel(): + self._reset_term.row_ids[ids] = torch.multinomial(probabilities, ids.numel(), replacement=True) + + cumulative_progress = self._progress_successes.sum().float() / self._attempts.sum().clamp_min(1) + total_success = self._final_successes.sum().float() / self._attempts.sum().clamp_min(1) + entropy = -(probabilities * probabilities.clamp_min(torch.finfo(probabilities.dtype).tiny).log()).sum() + entropy /= math.log(probabilities.numel()) + metrics: dict[str, torch.Tensor] = { + "batch_progress_rate": batch_progress, + "batch_success_rate": batch_success, + "batch_transfer_count": ( + command.transfer_counts[completed_ids].float().mean() + if completed_ids.numel() + else torch.zeros((), dtype=torch.float32, device=env.device) + ), + "batch_left_to_right_transfers": ( + command.direction_transfer_counts[completed_ids, 0].float().mean() + if completed_ids.numel() + else torch.zeros((), dtype=torch.float32, device=env.device) + ), + "batch_right_to_left_transfers": ( + command.direction_transfer_counts[completed_ids, 1].float().mean() + if completed_ids.numel() + else torch.zeros((), dtype=torch.float32, device=env.device) + ), + "deployment_probability": deployment_probability, + "row_coverage": row_coverage, + "overall_progress_rate": total_progress, + "cumulative_progress_rate": cumulative_progress, + "overall_success_rate": total_success, + "sampling_entropy": entropy, + } + for recipe in ConveyorResetRecipe: + mask = self._reset_term.recipe_ids == int(recipe) + recipe_attempts = self._attempts[mask].sum() + metrics[f"recipe_{recipe.name.lower()}_probability"] = probabilities[mask].sum() + recipe_history_size = history_size[mask].sum() + metrics[f"recipe_{recipe.name.lower()}_progress_rate"] = history_success_count[ + mask + ].sum() / recipe_history_size.clamp_min(1) + metrics[f"recipe_{recipe.name.lower()}_success_rate"] = self._final_successes[ + mask + ].sum().float() / recipe_attempts.clamp_min(1) + for side_id, side_name in ((0, "left_to_right"), (1, "right_to_left")): + mask = self._reset_term.source_side_ids == side_id + attempts = self._attempts[mask].sum() + direction_history_size = history_size[mask].sum() + metrics[f"direction_{side_name}_probability"] = probabilities[mask].sum() + metrics[f"direction_{side_name}_progress_rate"] = history_success_count[ + mask + ].sum() / direction_history_size.clamp_min(1) + metrics[f"direction_{side_name}_success_rate"] = self._final_successes[ + mask + ].sum().float() / attempts.clamp_min(1) + for recipe_name, variant_id, mask in self._diagnostic_variant_rows: + attempts = self._attempts[mask].sum() + variant_history_size = history_size[mask].sum() + prefix = f"recipe_{recipe_name}_variant_{variant_id}" + metrics[f"{prefix}_probability"] = probabilities[mask].sum() + metrics[f"{prefix}_progress_rate"] = history_success_count[mask].sum() / variant_history_size.clamp_min(1) + metrics[f"{prefix}_success_rate"] = self._final_successes[mask].sum().float() / attempts.clamp_min(1) + return metrics + + def get_state(self) -> dict[str, torch.Tensor]: + """Return curriculum evidence for checkpointing.""" + history_success_count = self._progress_monitor.success_buf.sum(dim=1).long() + return { + "attempts": self._attempts.clone(), + "progress_successes": self._progress_successes.clone(), + "final_successes": self._final_successes.clone(), + "progress_history": self._progress_monitor.success_buf.bool().clone(), + "history_pointer": self._progress_monitor.success_pointer.clone(), + "history_size": self._progress_monitor.success_size.clone(), + "history_success_count": history_success_count, + "rolling_progress_rates": self._progress_monitor.success_rate.clone(), + } + + def set_state(self, state: dict[str, torch.Tensor]) -> None: + """Restore curriculum evidence from a checkpoint.""" + targets = { + "attempts": self._attempts, + "progress_successes": self._progress_successes, + "final_successes": self._final_successes, + "progress_history": self._progress_monitor.success_buf, + "history_pointer": self._progress_monitor.success_pointer, + "history_size": self._progress_monitor.success_size, + "rolling_progress_rates": self._progress_monitor.success_rate, + } + for name, target in targets.items(): + if name not in state or state[name].shape != target.shape: + raise ValueError(f"Conveyor curriculum checkpoint has invalid '{name}'.") + if "history_success_count" not in state or state["history_success_count"].shape != self._attempts.shape: + raise ValueError("Conveyor curriculum checkpoint has invalid 'history_success_count'.") + history_len = self._progress_monitor.success_buf.shape[1] + if bool(torch.any((state["history_pointer"] < 0) | (state["history_pointer"] >= history_len))): + raise ValueError("Conveyor curriculum checkpoint has invalid history pointers.") + if bool(torch.any((state["history_size"] < 0) | (state["history_size"] > history_len))): + raise ValueError("Conveyor curriculum checkpoint has invalid history sizes.") + if bool( + torch.any((state["history_success_count"] < 0) | (state["history_success_count"] > state["history_size"])) + ): + raise ValueError("Conveyor curriculum checkpoint has invalid rolling success counts.") + for name, target in targets.items(): + target.copy_(state[name].to(device=target.device, dtype=target.dtype)) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/kinematics.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/kinematics.py new file mode 100644 index 000000000000..29ab17c9ce6f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/kinematics.py @@ -0,0 +1,71 @@ +# 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 + +"""Backend-independent Franka tool-state helpers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import math as math_utils + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +def _end_effector_cache_entry( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg, + body_name: str, + body_offset: tuple[float, float, float], +) -> tuple[Articulation, int, torch.Tensor]: + """Resolve and cache the Franka hand body and tool-frame offset.""" + robot: Articulation = env.scene[robot_cfg.name] + cache = getattr(env, "_conveyor_end_effector_cache", None) + if cache is None: + cache = {} + env._conveyor_end_effector_cache = cache + key = (robot_cfg.name, body_name, body_offset) + entry = cache.get(key) + if entry is None: + body_ids, _ = robot.find_bodies(body_name) + if len(body_ids) != 1: + raise ValueError(f"Expected one end-effector body matching '{body_name}', found {len(body_ids)}.") + entry = (body_ids[0], torch.tensor(body_offset, dtype=torch.float32, device=env.device)) + cache[key] = entry + return robot, entry[0], entry[1] + + +def end_effector_pose( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + body_name: str = "panda_hand", + body_offset: tuple[float, float, float] = (0.0, 0.0, 0.1034), +) -> tuple[torch.Tensor, torch.Tensor]: + """Return the Franka tool-center position [m] and orientation.""" + robot, body_id, offset = _end_effector_cache_entry(env, robot_cfg, body_name, body_offset) + orientation = robot.data.body_quat_w.torch[:, body_id] + position = robot.data.body_pos_w.torch[:, body_id] + position = position + math_utils.quat_apply(orientation, offset.expand(env.num_envs, -1)) + return position, orientation + + +def tool_velocity( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + body_name: str = "panda_hand", + body_offset: tuple[float, float, float] = (0.0, 0.0, 0.1034), +) -> torch.Tensor: + """Return tool-center linear and angular velocity [m/s, rad/s].""" + robot, body_id, offset = _end_effector_cache_entry(env, robot_cfg, body_name, body_offset) + orientation = robot.data.body_quat_w.torch[:, body_id] + body_velocity = robot.data.body_vel_w.torch[:, body_id] + offset_world = math_utils.quat_apply(orientation, offset.expand(env.num_envs, -1)) + linear_velocity = body_velocity[:, :3] + torch.linalg.cross(body_velocity[:, 3:], offset_world) + return torch.cat((linear_velocity, body_velocity[:, 3:]), dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/observations.py new file mode 100644 index 000000000000..8fb26f8d5460 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/observations.py @@ -0,0 +1,147 @@ +# 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 + +"""Task-conditioned observations for conveyor transfer.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import math as math_utils + +from ..conveyor_cube_pool import cube_values +from .kinematics import end_effector_pose, tool_velocity +from .reset_events import CUBE_COUNT, TRANSFER_X, side_inner_y + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +def _transfer_command(env: ManagerBasedRLEnv, command_name: str = "transfer"): + """Return the configured transfer command.""" + return env.command_manager.get_term(command_name) + + +def _cube_state(env: ManagerBasedRLEnv) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Stack cube world positions, orientations, and spatial velocities.""" + state = ( + cube_values(env, "root_pos_w"), + cube_values(env, "root_quat_w"), + cube_values(env, "root_vel_w"), + ) + adapter = getattr(env, "_adapt_policy_cube_state", None) + return state if adapter is None else adapter(*state) + + +def _active_cube_values(values: torch.Tensor, target_cube_ids: torch.Tensor) -> torch.Tensor: + """Gather one cube row for every vectorized environment.""" + shape = (values.shape[0], 1, *values.shape[2:]) + index = target_cube_ids.view(values.shape[0], 1, *([1] * (values.ndim - 2))).expand(shape) + return torch.gather(values, 1, index).squeeze(1) + + +def target_cube_one_hot(env: ManagerBasedRLEnv, command_name: str = "transfer") -> torch.Tensor: + """Encode which numbered cube the policy must transfer.""" + command = _transfer_command(env, command_name) + return torch.nn.functional.one_hot(command.target_cube_ids.long(), num_classes=CUBE_COUNT).float() + + +def target_side_one_hot(env: ManagerBasedRLEnv, command_name: str = "transfer") -> torch.Tensor: + """Encode the destination conveyor, opposite the reset source side.""" + command = _transfer_command(env, command_name) + return torch.nn.functional.one_hot(1 - command.source_side_ids.long(), num_classes=2).float() + + +def classify_cube_conveyors(local_positions: torch.Tensor, transit_half_width: float = 0.14) -> torch.Tensor: + """Classify positions as left conveyor, in transit, or right conveyor.""" + if local_positions.shape[-1] != 3: + raise ValueError("Cube positions must end in xyz coordinates.") + side_ids = torch.full(local_positions.shape[:-1], 1, dtype=torch.long, device=local_positions.device) + side_ids[local_positions[..., 1] > transit_half_width] = 0 + side_ids[local_positions[..., 1] < -transit_half_width] = 2 + return torch.nn.functional.one_hot(side_ids, num_classes=3).float() + + +def cube_conveyor_state(env: ManagerBasedRLEnv) -> torch.Tensor: + """Return left/transit/right one-hot state for every numbered cube.""" + positions, _, _ = _cube_state(env) + local_positions = positions - env.scene.env_origins.unsqueeze(1) + return classify_cube_conveyors(local_positions).flatten(start_dim=1) + + +def transfer_object_observation(env: ManagerBasedRLEnv) -> torch.Tensor: + """Describe all four cubes in stable identity slots. + + The observation contains local positions, tool-relative positions, local + up axes, and linear/angular velocities. The base task keeps fixed identities; + warehouse playback can refill remote slots from its physical parcel pool. + :func:`target_cube_one_hot` selects the active slot. + """ + positions, quaternions, velocities = _cube_state(env) + local_positions = positions - env.scene.env_origins.unsqueeze(1) + tool_position, _ = end_effector_pose(env) + tool_relative = positions - tool_position.unsqueeze(1) + rotations = math_utils.matrix_from_quat(quaternions.flatten(end_dim=1)).view(env.num_envs, CUBE_COUNT, 3, 3) + up_axes = rotations[..., 2] + return torch.cat( + ( + local_positions.flatten(start_dim=1), + tool_relative.flatten(start_dim=1), + up_axes.flatten(start_dim=1), + velocities.flatten(start_dim=1), + ), + dim=1, + ) + + +def active_transfer_features(env: ManagerBasedRLEnv, command_name: str = "transfer") -> torch.Tensor: + """Return active-cube and destination-relative position features [m].""" + command = _transfer_command(env, command_name) + positions, _, _ = _cube_state(env) + active_position = _active_cube_values(positions, command.target_cube_ids.long()) + local_active_position = active_position - env.scene.env_origins + tool_position, _ = end_effector_pose(env) + target_side_ids = 1 - command.source_side_ids.long() + target_position = torch.stack( + ( + torch.full_like(target_side_ids, TRANSFER_X, dtype=active_position.dtype), + side_inner_y(target_side_ids), + torch.full_like(target_side_ids, 0.06, dtype=active_position.dtype), + ), + dim=1, + ) + return torch.cat( + ( + local_active_position, + active_position - tool_position, + target_position - local_active_position, + ), + dim=1, + ) + + +def gripper_joint_positions( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot", joint_names=["panda_finger_joint[1-2]"]), +) -> torch.Tensor: + """Return the two Franka finger positions [m].""" + robot: Articulation = env.scene[robot_cfg.name] + return robot.data.joint_pos.torch[:, robot_cfg.joint_ids] + + +def end_effector_axes(env: ManagerBasedRLEnv) -> torch.Tensor: + """Return continuous tool-frame x and z axes.""" + _, orientation = end_effector_pose(env) + rotation = math_utils.matrix_from_quat(orientation) + return torch.cat((rotation[:, :, 0], rotation[:, :, 2]), dim=1) + + +def end_effector_velocity(env: ManagerBasedRLEnv) -> torch.Tensor: + """Return tool-center linear and angular velocity [m/s, rad/s].""" + return tool_velocity(env) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/reset_events.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/reset_events.py new file mode 100644 index 000000000000..5f6ff0dd9017 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/reset_events.py @@ -0,0 +1,545 @@ +# 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 + +"""Reset-state sampling for conveyor transfer.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from enum import IntEnum +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import EventTermCfg, ManagerTermBase + +from ..conveyor_geometry import ( + BELT_CENTER_X, + BELT_INNER_STRAIGHT_Y, + BELT_OUTER_STRAIGHT_Y, + BELT_TOP_Z, + CUBE_INNER_SLOT_X, + CUBE_OUTER_SLOT_X, +) + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +CUBE_COUNT = 4 +CUBE_SIZE = 0.04 +CUBE_REST_Z = BELT_TOP_Z + 0.5 * CUBE_SIZE +TRANSFER_X = 0.52 +LEFT_SIDE = 0 +RIGHT_SIDE = 1 +_BELT_CURRICULUM_FRACTIONS = (0.0, 0.15, 0.30, 0.45, 0.60, 0.80, 1.0) +_GRASP_CLOSURE_FRACTIONS = (0.25, 0.50, 0.75, 0.90, 1.0) +_OPEN_FINGER_POSITION = 0.040 +_CLOSED_FINGER_POSITION = 0.019 +BELT_DEPLOYMENT_VARIANT = len(_BELT_CURRICULUM_FRACTIONS) - 1 + + +class ConveyorResetRecipe(IntEnum): + """Reset phases ordered from easiest to complete task start.""" + + GOAL = 0 + PLACE = 1 + CARRY = 2 + LIFT = 3 + GRASP = 4 + PREGRASP = 5 + BELT = 6 + + +@dataclass(frozen=True) +class ConveyorResetRow: + """One physical reset state and transfer command.""" + + recipe: ConveyorResetRecipe + variant_id: int + target_cube_id: int + source_side_id: int + arm_positions: tuple[float, ...] + finger_position: float + held: bool + belt_range_fraction: float + + +_HOME_ARM = (0.0, -0.35, 0.0, -2.35, 0.0, 2.0, 0.78) + +# Constrained IK anchors target x=0.52 m with a downward-facing tool. Side 0 +# is the positive-y conveyor and side 1 is the negative-y conveyor. +_SOURCE_GRASP_ARM = ( + (0.5144946, 0.5427007, -0.0340481, -2.0731085, 0.0350549, 2.6153000, 1.2350098), + (-0.3400387, 0.5475255, -0.1330232, -2.0714802, 0.1368175, 2.6109921, 0.2080209), +) +_SOURCE_PREGRASP_ARM = ( + (0.5190744, 0.3742786, -0.0406565, -2.0872751, 0.0236238, 2.4611441, 1.2428656), + (-0.3220402, 0.3787346, -0.1588824, -2.0865274, 0.0928583, 2.4588233, 0.2380420), +) +_SOURCE_LIFT_ARM = ( + (0.5212608, 0.2372272, -0.0439391, -2.0514939, 0.0137043, 2.2884652, 1.2495379), + (-0.3132103, 0.2404478, -0.1720021, -2.0512362, 0.0540920, 2.2875542, 0.2640870), +) +_TARGET_PLACE_ARM = ( + # Source left, target right. + (-0.2921792, 0.4521663, -0.1853937, -2.0854943, 0.1398595, 2.5262992, 0.2062777), + # Source right, target left. + (0.4780173, 0.4442472, 0.0008774, -2.0873070, -0.0005178, 2.5315716, 1.2591928), +) +_CARRY_ARM = (0.0, 0.0137827, 0.0, -2.2661237, 0.0, 2.2798989, 0.78) + + +def _interpolate_arm( + start: tuple[float, ...], + end: tuple[float, ...], + fractions: tuple[float, ...], +) -> tuple[tuple[float, ...], ...]: + """Linearly interpolate validated joint-space anchors.""" + return tuple( + tuple( + start_value + fraction * (end_value - start_value) + for start_value, end_value in zip(start, end, strict=True) + ) + for fraction in fractions + ) + + +def _arm_position_variants( + recipe: ConveyorResetRecipe, + source_side_id: int, +) -> tuple[tuple[float, ...], ...]: + """Return dense reset states along the nominal transfer trajectory.""" + source_grasp = _SOURCE_GRASP_ARM[source_side_id] + source_pregrasp = _SOURCE_PREGRASP_ARM[source_side_id] + source_lift = _SOURCE_LIFT_ARM[source_side_id] + target_lift = _SOURCE_LIFT_ARM[1 - source_side_id] + if recipe == ConveyorResetRecipe.GOAL: + return (_SOURCE_PREGRASP_ARM[1 - source_side_id],) + if recipe == ConveyorResetRecipe.PLACE: + return _interpolate_arm(target_lift, _TARGET_PLACE_ARM[source_side_id], (0.25, 0.50, 0.75, 1.0)) + if recipe == ConveyorResetRecipe.CARRY: + return ( + *_interpolate_arm(source_lift, _CARRY_ARM, (0.33, 0.66, 1.0)), + *_interpolate_arm(_CARRY_ARM, target_lift, (0.33, 0.66, 1.0)), + ) + if recipe == ConveyorResetRecipe.LIFT: + return _interpolate_arm(source_grasp, source_lift, (0.25, 0.50, 0.75, 1.0)) + if recipe == ConveyorResetRecipe.GRASP: + return (source_grasp,) * len(_GRASP_CLOSURE_FRACTIONS) + if recipe == ConveyorResetRecipe.PREGRASP: + # Include the exact open-gripper acquisition pose. The previous last + # row stopped at 88% of the approach, leaving an approximately 1 cm + # gap between reset-driven approach learning and the already-held + # GRASP row. Dense near-contact rows make the physical close-and-lift + # transition learnable from the same sparse delivery objective. + return _interpolate_arm(source_pregrasp, source_grasp, (0.0, 0.50, 0.75, 0.92, 1.0)) + if recipe == ConveyorResetRecipe.BELT: + return _interpolate_arm(source_pregrasp, _HOME_ARM, _BELT_CURRICULUM_FRACTIONS) + raise ValueError(f"Unsupported reset recipe: {recipe}.") + + +def _finger_position_variants(recipe: ConveyorResetRecipe) -> tuple[float, ...]: + """Return finger positions paired with a recipe's arm variants [m].""" + variant_count = len(_arm_position_variants(recipe, LEFT_SIDE)) + if recipe == ConveyorResetRecipe.GRASP: + return tuple( + _OPEN_FINGER_POSITION + fraction * (_CLOSED_FINGER_POSITION - _OPEN_FINGER_POSITION) + for fraction in _GRASP_CLOSURE_FRACTIONS + ) + if recipe in (ConveyorResetRecipe.LIFT, ConveyorResetRecipe.CARRY, ConveyorResetRecipe.PLACE): + return (_CLOSED_FINGER_POSITION,) * variant_count + return (_OPEN_FINGER_POSITION,) * variant_count + + +def reset_variant_counts() -> tuple[int, ...]: + """Return the number of trajectory variants in each reset recipe.""" + return tuple(len(_arm_position_variants(recipe, LEFT_SIDE)) for recipe in ConveyorResetRecipe) + + +def build_reset_rows() -> tuple[ConveyorResetRow, ...]: + """Build the complete identity, direction, and dense-trajectory cross product.""" + return tuple( + ConveyorResetRow( + recipe=recipe, + variant_id=variant_id, + target_cube_id=cube_id, + source_side_id=source_side, + arm_positions=arm_positions, + finger_position=finger_position, + held=( + recipe in (ConveyorResetRecipe.LIFT, ConveyorResetRecipe.CARRY, ConveyorResetRecipe.PLACE) + or (recipe == ConveyorResetRecipe.GRASP and variant_id == len(_GRASP_CLOSURE_FRACTIONS) - 1) + ), + belt_range_fraction=_BELT_CURRICULUM_FRACTIONS[variant_id] if recipe == ConveyorResetRecipe.BELT else 0.0, + ) + for recipe in ConveyorResetRecipe + for cube_id in range(CUBE_COUNT) + for source_side in (LEFT_SIDE, RIGHT_SIDE) + for variant_id, (arm_positions, finger_position) in enumerate( + zip(_arm_position_variants(recipe, source_side), _finger_position_variants(recipe), strict=True) + ) + ) + + +_FRANKA_JOINT_ORIGINS = ( + ((0.0, 0.0, 0.333), (0.0, 0.0, 0.0)), + ((0.0, 0.0, 0.0), (-math.pi / 2.0, 0.0, 0.0)), + ((0.0, -0.316, 0.0), (math.pi / 2.0, 0.0, 0.0)), + ((0.0825, 0.0, 0.0), (math.pi / 2.0, 0.0, 0.0)), + ((-0.0825, 0.384, 0.0), (-math.pi / 2.0, 0.0, 0.0)), + ((0.0, 0.0, 0.0), (math.pi / 2.0, 0.0, 0.0)), + ((0.088, 0.0, 0.0), (math.pi / 2.0, 0.0, 0.0)), +) + + +def _rotation_matrix_from_rpy(roll: float, pitch: float, yaw: float, reference: torch.Tensor) -> torch.Tensor: + """Return a fixed XYZ roll-pitch-yaw rotation matrix.""" + cr, sr = math.cos(roll), math.sin(roll) + cp, sp = math.cos(pitch), math.sin(pitch) + cy, sy = math.cos(yaw), math.sin(yaw) + return reference.new_tensor( + ( + (cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr), + (sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr), + (-sp, cp * sr, cp * cr), + ) + ) + + +def franka_tool_position(joint_positions: torch.Tensor) -> torch.Tensor: + """Compute Panda tool-center positions [m] from seven joints [rad].""" + if joint_positions.shape[-1] != 7: + raise ValueError("Franka reset forward kinematics expects seven joint positions.") + shape = joint_positions.shape[:-1] + joints = joint_positions.reshape(-1, 7) + count = joints.shape[0] + rotation = torch.eye(3, dtype=joints.dtype, device=joints.device).expand(count, -1, -1).clone() + position = torch.zeros((count, 3), dtype=joints.dtype, device=joints.device) + reference = joints[0] if count else joint_positions.new_zeros(7) + for joint_id, (origin_position, origin_rpy) in enumerate(_FRANKA_JOINT_ORIGINS): + origin = joint_positions.new_tensor(origin_position).expand(count, -1) + position += torch.bmm(rotation, origin.unsqueeze(-1)).squeeze(-1) + rotation = torch.matmul(rotation, _rotation_matrix_from_rpy(*origin_rpy, reference=reference)) + angle = joints[:, joint_id] + cosine, sine = torch.cos(angle), torch.sin(angle) + zeros, ones = torch.zeros_like(angle), torch.ones_like(angle) + joint_rotation = torch.stack( + ( + torch.stack((cosine, -sine, zeros), dim=1), + torch.stack((sine, cosine, zeros), dim=1), + torch.stack((zeros, zeros, ones), dim=1), + ), + dim=1, + ) + rotation = torch.bmm(rotation, joint_rotation) + tool_offset = joint_positions.new_tensor((0.0, 0.0, 0.2104)).expand(count, -1) + position += torch.bmm(rotation, tool_offset.unsqueeze(-1)).squeeze(-1) + return position.reshape(*shape, 3) + + +def side_inner_y(side_ids: torch.Tensor) -> torch.Tensor: + """Return the reachable inner-straight y coordinate [m] for each side.""" + return torch.where(side_ids == LEFT_SIDE, BELT_INNER_STRAIGHT_Y, -BELT_INNER_STRAIGHT_Y) + + +def _balanced_cube_slots( + target_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Assign one cube to every inner/outer racetrack run. + + Slots are ordered ``left inner``, ``left outer``, ``right inner``, and + ``right outer``. The commanded cube swaps with the canonical occupant of + its source-side inner slot, keeping it reachable without duplicating or + emptying any deployment slot. + """ + if target_cube_ids.ndim != 1 or source_side_ids.shape != target_cube_ids.shape: + raise ValueError("Target cube and source-side ids must be matching vectors.") + if bool(torch.any((target_cube_ids < 0) | (target_cube_ids >= CUBE_COUNT))): + raise ValueError("Target cube ids are out of range.") + if bool(torch.any((source_side_ids != LEFT_SIDE) & (source_side_ids != RIGHT_SIDE))): + raise ValueError("Source-side ids must be 0 (left) or 1 (right).") + + count = target_cube_ids.numel() + cube_slots = torch.arange(CUBE_COUNT, device=target_cube_ids.device).expand(count, -1).clone() + source_inner_slots = 2 * source_side_ids + displaced_cube_ids = source_inner_slots + target_original_slots = target_cube_ids + cube_slots.scatter_(1, target_cube_ids.unsqueeze(1), source_inner_slots.unsqueeze(1)) + cube_slots.scatter_(1, displaced_cube_ids.unsqueeze(1), target_original_slots.unsqueeze(1)) + + cube_sides = torch.div(cube_slots, 2, rounding_mode="floor") + on_outer_run = torch.remainder(cube_slots, 2).bool() + cube_x = torch.where(on_outer_run, CUBE_OUTER_SLOT_X, CUBE_INNER_SLOT_X) + y_magnitude = torch.where(on_outer_run, BELT_OUTER_STRAIGHT_Y, BELT_INNER_STRAIGHT_Y) + cube_y = torch.where(cube_sides == LEFT_SIDE, y_magnitude, -y_magnitude) + return cube_slots, cube_sides, cube_x, cube_y + + +def _sample_collision_free_active_x( + base_x: torch.Tensor, + cube_sides: torch.Tensor, + target_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, + lower: torch.Tensor, + upper: torch.Tensor, + minimum_separation: float = 0.055, + attempts: int = 16, +) -> torch.Tensor: + """Sample active-cube positions without overlapping an inactive cube.""" + if base_x.ndim != 2 or base_x.shape[1] != CUBE_COUNT or cube_sides.shape != base_x.shape: + raise ValueError("base_x and cube_sides must have shape (N, CUBE_COUNT).") + count = base_x.shape[0] + expected_vector_shape = (count,) + if any(value.shape != expected_vector_shape for value in (target_cube_ids, source_side_ids, lower, upper)): + raise ValueError("Active-cube sampling controls must have shape (N,).") + if minimum_separation <= 0.0 or attempts < 1: + raise ValueError("Invalid collision-free active-cube sampling parameters.") + + cube_ids = torch.arange(CUBE_COUNT, device=base_x.device).expand(count, -1) + inactive_on_source = (cube_sides == source_side_ids.unsqueeze(1)) & (cube_ids != target_cube_ids.unsqueeze(1)) + sampled = lower + torch.rand_like(lower) * (upper - lower) + for _ in range(attempts): + conflicts = torch.any( + (torch.abs(sampled.unsqueeze(1) - base_x) < minimum_separation) & inactive_on_source, dim=1 + ) + replacement = lower + torch.rand_like(lower) * (upper - lower) + sampled = torch.where(conflicts, replacement, sampled) + + conflicts = torch.any((torch.abs(sampled.unsqueeze(1) - base_x) < minimum_separation) & inactive_on_source, dim=1) + fallback = torch.full_like(sampled, TRANSFER_X) + return torch.where(conflicts, fallback, sampled) + + +class ConveyorResetStateTable(ManagerTermBase): + """Restore validated transfer states spanning released goal to moving start.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self._rows = build_reset_rows() + self.recipe_ids = torch.tensor([row.recipe for row in self._rows], dtype=torch.long, device=env.device) + self.target_cube_ids = torch.tensor( + [row.target_cube_id for row in self._rows], dtype=torch.long, device=env.device + ) + self.variant_ids = torch.tensor([row.variant_id for row in self._rows], dtype=torch.long, device=env.device) + self.source_side_ids = torch.tensor( + [row.source_side_id for row in self._rows], dtype=torch.long, device=env.device + ) + self._arm_positions = torch.tensor( + [row.arm_positions for row in self._rows], dtype=torch.float32, device=env.device + ) + self._finger_positions = torch.tensor( + [row.finger_position for row in self._rows], dtype=torch.float32, device=env.device + ) + self.held_rows = torch.tensor([row.held for row in self._rows], dtype=torch.bool, device=env.device) + self._belt_range_fractions = torch.tensor( + [row.belt_range_fraction for row in self._rows], dtype=torch.float32, device=env.device + ) + self._robot: Articulation = env.scene["robot"] + self._cubes: tuple[RigidObject, ...] = tuple(env.scene[f"cube_{cube_id}"] for cube_id in range(CUBE_COUNT)) + self._arm_joint_ids = self._robot.find_joints("panda_joint[1-7]", preserve_order=True)[0] + self._finger_joint_ids = self._robot.find_joints("panda_finger_joint[1-2]", preserve_order=True)[0] + if len(self._arm_joint_ids) != 7 or len(self._finger_joint_ids) != 2: + raise ValueError("Conveyor transfer requires seven Panda arm joints and two finger joints.") + self.row_ids = torch.randint(self.row_count, (env.num_envs,), dtype=torch.long, device=env.device) + + @property + def row_count(self) -> int: + """Number of immutable physical reset rows.""" + return len(self._rows) + + @property + def recipe_names(self) -> tuple[str, ...]: + """Stable reset recipe labels.""" + return tuple(recipe.name.lower() for recipe in ConveyorResetRecipe) + + def _filtered_rows( + self, + fixed_recipe: int | None, + fixed_variant_id: int | None, + fixed_target_cube_id: int | None, + fixed_source_side_id: int | None, + ) -> torch.Tensor: + """Return rows matching optional deterministic evaluation controls.""" + mask = torch.ones(self.row_count, dtype=torch.bool, device=self.device) + if fixed_recipe is not None: + if not 0 <= fixed_recipe < len(ConveyorResetRecipe): + raise ValueError(f"fixed_recipe must lie in [0, {len(ConveyorResetRecipe) - 1}].") + mask &= self.recipe_ids == fixed_recipe + if fixed_variant_id is not None: + if fixed_variant_id < 0: + raise ValueError("fixed_variant_id must be non-negative.") + mask &= self.variant_ids == fixed_variant_id + if fixed_target_cube_id is not None: + if not 0 <= fixed_target_cube_id < CUBE_COUNT: + raise ValueError(f"fixed_target_cube_id must lie in [0, {CUBE_COUNT - 1}].") + mask &= self.target_cube_ids == fixed_target_cube_id + if fixed_source_side_id is not None: + if fixed_source_side_id not in (LEFT_SIDE, RIGHT_SIDE): + raise ValueError("fixed_source_side_id must be 0 (left) or 1 (right).") + mask &= self.source_side_ids == fixed_source_side_id + return torch.nonzero(mask, as_tuple=False).flatten() + + def __call__( + self, + env: ManagerBasedRLEnv, + env_ids: torch.Tensor, + fixed_recipe: int | None = None, + fixed_variant_id: int | None = None, + fixed_target_cube_id: int | None = None, + fixed_source_side_id: int | None = None, + belt_start_x_range: tuple[float, float] = (0.30, 0.82), + cube_position_noise: float = 0.015, + arm_joint_noise: float = 0.015, + ) -> None: + """Write sampled robot and four-cube states directly into simulation.""" + if env_ids is None or env_ids.numel() == 0: + return + if belt_start_x_range[0] >= belt_start_x_range[1]: + raise ValueError("belt_start_x_range must be strictly increasing.") + if cube_position_noise < 0.0 or arm_joint_noise < 0.0: + raise ValueError("Reset randomization ranges must be non-negative.") + + if ( + fixed_recipe is None + and fixed_variant_id is None + and fixed_target_cube_id is None + and fixed_source_side_id is None + ): + row_ids = self.row_ids[env_ids] + else: + candidates = self._filtered_rows( + fixed_recipe, + fixed_variant_id, + fixed_target_cube_id, + fixed_source_side_id, + ) + if candidates.numel() == 0: + raise RuntimeError("No conveyor reset rows match the fixed reset controls.") + row_ids = candidates[torch.randint(candidates.numel(), (env_ids.numel(),), device=self.device)] + self.row_ids[env_ids] = row_ids + + recipes = self.recipe_ids[row_ids] + target_cube_ids = self.target_cube_ids[row_ids] + source_side_ids = self.source_side_ids[row_ids] + held_rows = self.held_rows[row_ids] + + arm_positions = self._arm_positions[row_ids].clone() + if arm_joint_noise > 0.0: + noise = (2.0 * torch.rand_like(arm_positions) - 1.0) * arm_joint_noise + # Preserve the exact approach, closure, and held-object manifold. + # Only deployment-like BELT starts receive robot randomization. + noise[recipes != int(ConveyorResetRecipe.BELT)] = 0.0 + arm_positions += noise + joint_positions = self._robot.data.default_joint_pos.torch[env_ids].clone() + joint_velocities = torch.zeros_like(joint_positions) + joint_positions[:, self._arm_joint_ids] = arm_positions + finger_positions = self._finger_positions[row_ids].to(dtype=joint_positions.dtype).unsqueeze(1).expand(-1, 2) + joint_positions[:, self._finger_joint_ids] = finger_positions + self._robot.set_joint_position_target_index(target=joint_positions, env_ids=env_ids) + self._robot.set_joint_velocity_target_index(target=joint_velocities, env_ids=env_ids) + self._robot.write_joint_position_to_sim_index(position=joint_positions, env_ids=env_ids) + self._robot.write_joint_velocity_to_sim_index(velocity=joint_velocities, env_ids=env_ids) + + count = env_ids.numel() + cube_slots, cube_sides, base_x, cube_y = _balanced_cube_slots(target_cube_ids, source_side_ids) + base_x = base_x.to(dtype=arm_positions.dtype) + cube_y = cube_y.to(dtype=arm_positions.dtype) + if cube_position_noise > 0.0: + base_x += (2.0 * torch.rand_like(base_x) - 1.0) * cube_position_noise + + active_lower = torch.full((count,), TRANSFER_X, dtype=arm_positions.dtype, device=self.device) + active_upper = active_lower.clone() + belt_rows = recipes == int(ConveyorResetRecipe.BELT) + if bool(torch.any(belt_rows)): + range_fraction = self._belt_range_fractions[row_ids] + range_lower = TRANSFER_X + range_fraction * (belt_start_x_range[0] - TRANSFER_X) + range_upper = TRANSFER_X + range_fraction * (belt_start_x_range[1] - TRANSFER_X) + active_lower[belt_rows] = range_lower[belt_rows] + active_upper[belt_rows] = range_upper[belt_rows] + active_x = _sample_collision_free_active_x( + base_x, + cube_sides, + target_cube_ids, + source_side_ids, + active_lower, + active_upper, + ) + + # On a full deployment reset, mirror the other cube on the source + # conveyor across the racetrack center. Opposite straight runs then + # differ by exactly half a lap even when the active start is sampled. + source_outer_slots = 2 * source_side_ids + 1 + source_outer_cube = cube_slots == source_outer_slots.unsqueeze(1) + mirrored_outer_x = (2.0 * BELT_CENTER_X - active_x).unsqueeze(1) + base_x = torch.where(belt_rows.unsqueeze(1) & source_outer_cube, mirrored_outer_x, base_x) + cube_positions = torch.stack( + (base_x, cube_y, torch.full_like(base_x, CUBE_REST_Z)), + dim=2, + ) + + active_positions = torch.stack( + (active_x, side_inner_y(source_side_ids), torch.full_like(active_x, CUBE_REST_Z)), + dim=1, + ) + goal_rows = recipes == int(ConveyorResetRecipe.GOAL) + active_positions[goal_rows, 1] = side_inner_y(1 - source_side_ids[goal_rows]) + if bool(torch.any(held_rows)): + active_positions[held_rows] = franka_tool_position(arm_positions[held_rows]) + cube_positions.scatter_( + 1, + target_cube_ids.view(-1, 1, 1).expand(-1, 1, 3), + active_positions.unsqueeze(1), + ) + + identity_quaternion = arm_positions.new_tensor((0.0, 0.0, 0.0, 1.0)).expand(count, -1) + for cube_id, cube in enumerate(self._cubes): + root_pose = cube.data.default_root_pose.torch[env_ids].clone() + root_pose[:, :3] = cube_positions[:, cube_id] + env.scene.env_origins[env_ids] + root_pose[:, 3:7] = identity_quaternion + root_velocity = torch.zeros((count, 6), dtype=root_pose.dtype, device=self.device) + cube.write_root_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) + cube.write_root_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Keep the immutable reset table across environment resets.""" + + +def select_next_transfer_cube( + cube_positions: torch.Tensor, + current_cube_ids: torch.Tensor, + source_side_ids: torch.Tensor, + transit_half_width: float = 0.14, +) -> torch.Tensor: + """Sample the next numbered cube already located on each source belt. + + A different eligible cube is sampled uniformly. The just-placed cube is + the fallback, so a valid goal remains available even when it is temporarily + the only parcel on that conveyor. + """ + if cube_positions.ndim != 3 or cube_positions.shape[1:] != (CUBE_COUNT, 3): + raise ValueError(f"cube_positions must have shape (N, {CUBE_COUNT}, 3).") + count = cube_positions.shape[0] + if current_cube_ids.shape != (count,) or source_side_ids.shape != (count,): + raise ValueError("Current cube and source-side ids must match the position batch.") + if transit_half_width <= 0.0: + raise ValueError("transit_half_width must be positive.") + + cube_ids = torch.arange(CUBE_COUNT, device=cube_positions.device).expand(count, -1) + on_left = cube_positions[:, :, 1] > transit_half_width + on_right = cube_positions[:, :, 1] < -transit_half_width + candidates = torch.where(source_side_ids.unsqueeze(1) == LEFT_SIDE, on_left, on_right) + alternatives = candidates & (cube_ids != current_cube_ids.unsqueeze(1)) + candidates = torch.where(torch.any(alternatives, dim=1, keepdim=True), alternatives, candidates) + + has_candidates = torch.any(candidates, dim=1) + fallback = torch.zeros_like(candidates) + fallback.scatter_(1, current_cube_ids.unsqueeze(1), True) + candidates = torch.where(has_candidates.unsqueeze(1), candidates, fallback) + return torch.multinomial(candidates.float(), 1).squeeze(1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/rewards.py new file mode 100644 index 000000000000..1c161562eaac --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/rewards.py @@ -0,0 +1,156 @@ +# 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 + +"""Progress utilities and sparse completion rewards for conveyor transfer.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import SceneEntityCfg + +from ..conveyor_cube_pool import cube_values +from .kinematics import end_effector_pose +from .reset_events import CUBE_REST_Z, TRANSFER_X, side_inner_y + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + from .commands import ConveyorTransferCommand + + +def transfer_potential( + cube_positions: torch.Tensor, + tool_positions: torch.Tensor, + finger_positions: torch.Tensor, + source_side_ids: torch.Tensor, +) -> torch.Tensor: + """Return a monotonic pickup-to-release shaping potential.""" + source_y = side_inner_y(source_side_ids) + target_y = side_inner_y(1 - source_side_ids) + target_position = torch.stack( + ( + torch.full_like(source_y, TRANSFER_X), + target_y, + torch.full_like(source_y, CUBE_REST_Z), + ), + dim=1, + ) + tool_distance = torch.linalg.vector_norm(tool_positions - cube_positions, dim=1) + reach = torch.exp(-12.0 * tool_distance) + gripper_closure = torch.clamp((0.04 - torch.amin(finger_positions, dim=1)) / 0.021, min=0.0, max=1.0) + grasp = gripper_closure * torch.exp(-25.0 * tool_distance) + lift = torch.clamp((cube_positions[:, 2] - CUBE_REST_Z) / 0.14, min=0.0, max=1.0) + direction_denominator = (target_y - source_y).clamp(min=-1.0, max=1.0) + crossing = (cube_positions[:, 1] - source_y) / direction_denominator + crossing = torch.clamp(crossing, min=0.0, max=1.0) + target_distance = torch.linalg.vector_norm(cube_positions - target_position, dim=1) + target = torch.exp(-14.0 * target_distance) + transport = crossing * torch.maximum(torch.clamp(2.0 * lift, max=1.0), target) + released = (torch.amin(finger_positions, dim=1) > 0.027).float() * target + return 0.5 * reach + 0.75 * grasp + 1.25 * lift + 2.0 * transport + 2.0 * target + released + + +def current_transfer_potential( + env: ManagerBasedRLEnv, + command_name: str = "transfer", + command: ConveyorTransferCommand | None = None, +) -> torch.Tensor: + """Gather current task state and evaluate the shaping potential.""" + if command is None: + command = env.command_manager.get_term(command_name) + positions = cube_values(env, "root_pos_w") + index = command.target_cube_ids.view(env.num_envs, 1, 1).expand(-1, 1, 3) + active_position = torch.gather(positions, 1, index).squeeze(1) - env.scene.env_origins + tool_position, _ = end_effector_pose(env) + tool_position = tool_position - env.scene.env_origins + robot: Articulation = env.scene["robot"] + finger_ids, _ = robot.find_joints("panda_finger_joint[1-2]", preserve_order=True) + finger_positions = robot.data.joint_pos.torch[:, finger_ids] + return transfer_potential(active_position, tool_position, finger_positions, command.source_side_ids) + + +def transfer_success_reward( + env: ManagerBasedRLEnv, + command_name: str = "transfer", +) -> torch.Tensor: + """Return one on each stable transfer-completion transition.""" + command = env.command_manager.get_term(command_name) + command.evaluate() + return command.new_success.float() + + +def terminal_failure( + env: ManagerBasedRLEnv, + command_name: str = "transfer", +) -> torch.Tensor: + """Return one for non-timeout terminal failures.""" + command = env.command_manager.get_term(command_name) + command.evaluate() + return (env.reset_terminated & ~command.pending_success).float() + + +def action_term_l2(env: ManagerBasedRLEnv, action_name: str) -> torch.Tensor: + """Penalize one named raw action term.""" + action = env.action_manager.get_term(action_name).raw_actions + return torch.sum(torch.square(action), dim=1) + + +def finite_action_rate_l2( + env: ManagerBasedRLEnv, + action_names: tuple[str, ...] = ("arm_action", "gripper_action"), +) -> torch.Tensor: + """Penalize changes between the finite policy commands accepted by action terms.""" + if not action_names: + raise ValueError("At least one action term is required for the action-rate reward.") + terms = tuple(env.action_manager.get_term(name) for name in action_names) + action = torch.cat(tuple(term.raw_actions for term in terms), dim=1) + previous_action = torch.cat(tuple(term.previous_actions for term in terms), dim=1) + return torch.sum(torch.square(action - previous_action), dim=1) + + +def physical_cube_acquisition_mask( + env: ManagerBasedRLEnv, + command_name: str = "transfer", + command: ConveyorTransferCommand | None = None, + minimum_lift: float = 0.025, + maximum_tool_distance: float = 0.075, + maximum_finger_position: float = 0.030, +) -> torch.Tensor: + """Return physically closed, lifted, tool-local commanded-cube grasps.""" + if minimum_lift <= 0.0 or maximum_tool_distance <= 0.0 or maximum_finger_position <= 0.0: + raise ValueError("Physical acquisition thresholds must be positive.") + if command is None: + command = env.command_manager.get_term(command_name) + positions = cube_values(env, "root_pos_w") + index = command.target_cube_ids.view(env.num_envs, 1, 1).expand(-1, 1, 3) + active_position = torch.gather(positions, 1, index).squeeze(1) + tool_position, _ = end_effector_pose(env) + robot: Articulation = env.scene["robot"] + finger_ids, _ = robot.find_joints("panda_finger_joint[1-2]", preserve_order=True) + finger_positions = robot.data.joint_pos.torch[:, finger_ids] + local_cube_z = active_position[:, 2] - env.scene.env_origins[:, 2] + return ( + (local_cube_z >= CUBE_REST_Z + minimum_lift) + & (torch.linalg.vector_norm(active_position - tool_position, dim=1) <= maximum_tool_distance) + & (torch.amax(finger_positions, dim=1) <= maximum_finger_position) + ) + + +def finite_joint_velocity_l2( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot", joint_names=["panda_joint[1-7]"]), + maximum_velocity: float = 3.0, +) -> torch.Tensor: + """Penalize bounded arm velocity while sanitizing divergent states.""" + if maximum_velocity <= 0.0: + raise ValueError("maximum_velocity must be positive.") + robot: Articulation = env.scene[asset_cfg.name] + velocity = robot.data.joint_vel.torch[:, asset_cfg.joint_ids] + velocity = torch.nan_to_num(velocity, nan=0.0, posinf=maximum_velocity, neginf=-maximum_velocity) + return torch.sum(torch.square(torch.clamp(velocity, -maximum_velocity, maximum_velocity)), dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/sorting.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/sorting.py new file mode 100644 index 000000000000..7cb92cea07d5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/sorting.py @@ -0,0 +1,128 @@ +# 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 + +"""Parcel-class dispatch for the warehouse's unchanged four-slot transfer policy.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.utils.configclass import configclass + +from ..conveyor_cube_pool import cube_values +from .commands import ConveyorTransferCommand, ConveyorTransferCommandCfg +from .rewards import physical_cube_acquisition_mask + +if TYPE_CHECKING: + from ..conveyor_franka_warehouse_env import ConveyorFrankaWarehouseEnv + + +class ConveyorSortCommand(ConveyorTransferCommand): + """Transfer misplaced cartons and leave correctly sorted inventory circulating. + + Physical parcel IDs carry immutable destination classes. The dispatcher presents + one wrong-lane arrival through the pretrained policy's existing cube/side command; + color is a visual class label, not an additional policy observation. + """ + + cfg: ConveyorSortCommandCfg + + def __init__(self, cfg: ConveyorSortCommandCfg, env: ConveyorFrankaWarehouseEnv) -> None: + super().__init__(cfg, env) + if not 0.14 <= cfg.pickup_x_range[0] < cfg.pickup_x_range[1] <= 1.02: + raise ValueError("The sorting pickup window must lie within the original working straight [0.14, 1.02] m.") + if len(cfg.parcel_destinations) != len(env.conveyor_cube_pool.assets) or set(cfg.parcel_destinations) != {0, 1}: + raise ValueError("Parcel destinations must match the physical pool and include both conveyor IDs, 0 and 1.") + if len(cfg.parcel_colors) != len(cfg.parcel_destinations): + raise ValueError("Each physical parcel requires a color and a destination.") + for color in set(cfg.parcel_colors): + destinations = { + side for shade, side in zip(cfg.parcel_colors, cfg.parcel_destinations, strict=True) if shade == color + } + if len(destinations) != 1: + raise ValueError(f"All {color} parcels must share one destination conveyor.") + self.parcel_destinations = torch.tensor(cfg.parcel_destinations, device=self.device) + self.has_target = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + self.metrics["sorted_parcels"] = torch.zeros(self.num_envs, device=self.device) + self.metrics["batch_complete"] = torch.zeros(self.num_envs, device=self.device) + + def evaluate(self) -> None: + """Credit stable placements only while the dispatcher owns an active transfer.""" + self.new_success[~self.has_target] = False + self.is_success[~self.has_target] = False + self._last_evaluation_steps[~self.has_target] = self._env.episode_length_buf[~self.has_target] + super().evaluate() + + def _resample_command(self, env_ids: Sequence[int]) -> None: + # Reset metadata remains compatible with the shared reward and observation terms. + if self._resampling_from_reset: + super()._resample_command(env_ids) + self.has_target[env_ids] = False + self.held_cube_ids[env_ids] = -1 + self.pending_success[env_ids] = False + + def _update_command(self) -> None: + """Publish slot assignments before the manager computes the next policy observation.""" + env = self._env + pool = env.conveyor_cube_pool + positions = cube_values(env, "root_pos_w", all_cubes=True) - env.scene.env_origins[:, None] + local = env._in_workcell(positions) + rows = torch.arange(self.num_envs, device=self.device) + physical_target = pool.slot_ids[rows, self.target_cube_ids] + held = physical_cube_acquisition_mask(env, command=self) & self.has_target + self.has_target &= ~self.pending_success & (local[rows, physical_target] | held) + self.pending_success.zero_() + wrong_lane = (positions[..., 1] < 0).long() != self.parcel_destinations + candidates = ( + wrong_lane + & (positions[..., 0] > self.cfg.pickup_x_range[0]) + & (positions[..., 0] < self.cfg.pickup_x_range[1]) + & (positions[..., 1].abs() > 0.20) + & (positions[..., 1].abs() < 0.36) + & (positions[..., 2] > 0.04) + & (positions[..., 2] < 0.10) + ) + pool.refresh(positions, local, candidates, self.target_cube_ids, self.has_target) + available = candidates.gather(1, pool.slot_ids) + eligible = available.any(dim=1) & ~self.has_target + choices = torch.where(available, positions[..., 0].gather(1, pool.slot_ids), -torch.inf).argmax(dim=1) + for slot in range(4): + env_ids = torch.where(eligible & (choices == slot))[0] + if env_ids.numel(): + self.set_goal(slot, env_ids) + self.has_target[env_ids] = True + self.command_counter[env_ids] += 1 + # A class is counted only after leaving the elevated supply and landing on a loop. + velocities = cube_values(env, "root_lin_vel_w", all_cubes=True) + on_loop = (positions[..., 2] > 0.04) & (positions[..., 2] < 0.20) & (velocities[..., 2].abs() < 0.15) + active = torch.zeros_like(on_loop).scatter_( + 1, pool.slot_ids[rows, self.target_cube_ids, None], self.has_target[:, None] + ) + on_loop &= ~active + sorted_parcels = (~wrong_lane & on_loop).sum(dim=1) + self.metrics["sorted_parcels"].copy_(sorted_parcels) + self.metrics["batch_complete"].copy_(sorted_parcels == len(pool.assets)) + + +@configclass +class ConveyorSortCommandCfg(ConveyorTransferCommandCfg): + """Batch sorting with fixed physical classes and checkpoint-compatible commands.""" + + class_type: type[ConveyorSortCommand] | str = "{DIR}.sorting:ConveyorSortCommand" + + parcel_destinations: tuple[int, ...] = (0, 1) * 12 + """Destination per physical parcel: 0 is the positive-Y loop, 1 the negative-Y loop.""" + + parcel_colors: tuple[str, ...] = ("blue", "orange", "green", "purple") * 6 + """Authored color per physical parcel; blue, orange, green, and purple are available.""" + + randomize_arrivals: bool = True + """Shuffle physical parcel identities across authored start positions at each reset.""" + + pickup_x_range: tuple[float, float] = (0.35, 1.02) + """Longitudinal arrival window in policy workspace coordinates [m].""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/terminations.py new file mode 100644 index 000000000000..9f1f680319bb --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/conveyor_franka/mdp/terminations.py @@ -0,0 +1,113 @@ +# 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 + +"""Actual episode termination and truncation terms for conveyor transfer.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import SceneEntityCfg + +from ..conveyor_cube_pool import cube_values +from ..conveyor_geometry import ( + BELT_CENTER_X, + BELT_HALF_STRAIGHT, + BELT_TURN_RADIUS, + BELT_WIDTH, + GUARD_THICKNESS, +) +from .reset_events import CUBE_SIZE + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +_TRACK_X_CLEARANCE = BELT_TURN_RADIUS + 0.5 * BELT_WIDTH + GUARD_THICKNESS + CUBE_SIZE + + +def invalid_action( + env: ManagerBasedRLEnv, + action_names: Sequence[str] = ("arm_action", "gripper_action"), +) -> torch.Tensor: + """Terminate environments whose latest policy action contained a non-finite value. + + Args: + env: Manager-based conveyor environment. + action_names: Action terms exposing an ``invalid_actions`` Boolean tensor. + + Returns: + Per-environment invalid-action mask. + """ + if not action_names: + raise ValueError("At least one action term is required for invalid-action termination.") + masks = tuple(env.action_manager.get_term(name).invalid_actions for name in action_names) + return torch.stack(masks, dim=0).any(dim=0) + + +def subgoal_time_out( + env: ManagerBasedRLEnv, + timeout_s: float = 20.0, + command_name: str = "transfer", +) -> torch.Tensor: + """Truncate environments that make no transfer within one subgoal timeout [s].""" + if timeout_s <= 0.0: + raise ValueError("timeout_s must be positive.") + command = env.command_manager.get_term(command_name) + command.evaluate() + timeout_steps = math.ceil(timeout_s / env.step_dt) + elapsed = env.episode_length_buf - command.subgoal_start_steps + return (elapsed >= timeout_steps) & ~command.pending_success + + +def transfer_sequence_time_out( + env: ManagerBasedRLEnv, + maximum_transfers: int = 8, + command_name: str = "transfer", +) -> torch.Tensor: + """Truncate completed sequences so reset-state coverage remains fresh.""" + if maximum_transfers < 1: + raise ValueError("maximum_transfers must be positive.") + command = env.command_manager.get_term(command_name) + command.evaluate() + return command.transfer_counts >= maximum_transfers + + +def cube_out_of_workspace( + env: ManagerBasedRLEnv, + minimum: tuple[float, float, float] = ( + BELT_CENTER_X - BELT_HALF_STRAIGHT - _TRACK_X_CLEARANCE, + -1.05, + -0.05, + ), + maximum: tuple[float, float, float] = ( + BELT_CENTER_X + BELT_HALF_STRAIGHT + _TRACK_X_CLEARANCE, + 1.05, + 0.80, + ), +) -> torch.Tensor: + """Terminate when any cube leaves the complete guarded racetrack workspace.""" + positions = cube_values(env, "root_pos_w", all_cubes=True) + positions -= env.scene.env_origins.unsqueeze(1) + lower = positions.new_tensor(minimum) + upper = positions.new_tensor(maximum) + return torch.any((positions < lower) | (positions > upper), dim=(1, 2)) + + +def nonfinite_scene_state( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Terminate environments containing nonfinite robot or cube state.""" + robot: Articulation = env.scene[robot_cfg.name] + invalid = ~torch.all(torch.isfinite(robot.data.joint_pos.torch), dim=1) + invalid |= ~torch.all(torch.isfinite(robot.data.joint_vel.torch), dim=1) + invalid |= ~torch.isfinite(cube_values(env, "root_state_w", all_cubes=True)).all(dim=(1, 2)) + return invalid diff --git a/source/isaaclab_tasks/pyproject.toml b/source/isaaclab_tasks/pyproject.toml index 740a89dbaa45..7922ff174755 100644 --- a/source/isaaclab_tasks/pyproject.toml +++ b/source/isaaclab_tasks/pyproject.toml @@ -20,11 +20,15 @@ requires-python = ">=3.12" dependencies = [ "isaaclab", "isaaclab_assets", + "isaaclab_newton", + "isaaclab_physx", ] [tool.uv.sources] isaaclab = { path = "../isaaclab", editable = true } isaaclab_assets = { path = "../isaaclab_assets", editable = true } +isaaclab_newton = { path = "../isaaclab_newton", editable = true } +isaaclab_physx = { path = "../isaaclab_physx", editable = true } [project.urls] Homepage = "https://github.com/isaac-sim/IsaacLab" @@ -39,3 +43,4 @@ include = ["isaaclab_tasks", "isaaclab_tasks.*"] [tool.setuptools.package-data] "*" = ["*.pyi"] +"isaaclab_tasks.contrib.conveyor_franka" = ["assets/*.usda", "assets/*.usd"] diff --git a/source/isaaclab_tasks/test/contrib/test_contrib_environments_kit.py b/source/isaaclab_tasks/test/contrib/test_contrib_environments_kit.py index 6372b740647b..ee50f506f538 100644 --- a/source/isaaclab_tasks/test/contrib/test_contrib_environments_kit.py +++ b/source/isaaclab_tasks/test/contrib/test_contrib_environments_kit.py @@ -30,4 +30,4 @@ @pytest.mark.parametrize("task_name", contrib_environment_params("kit")) def test_contrib_environments_kit(task_name): - _run_environments(task_name, device="cuda", num_envs=num_envs(task_name)) + _run_environments(task_name, device=None, num_envs=num_envs(task_name)) diff --git a/source/isaaclab_tasks/test/contrib/test_conveyor_franka_asset_cfg.py b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_asset_cfg.py new file mode 100644 index 000000000000..8ea26a5fa60b --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_asset_cfg.py @@ -0,0 +1,615 @@ +# 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 + +"""Import-light checks for the Digital Twin conveyor playback scene.""" + +import math +import subprocess +import sys +from types import SimpleNamespace + +import gymnasium as gym +import pytest + +import isaaclab.sim as sim_utils + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_asset_env_cfg import ( + _PRESENTATION_ASSETS, + ConveyorFrankaA09A12EnvCfg, + _make_usd_subtree_visual_only, + _presentation_layer, +) +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env_cfg import ConveyorFrankaEnvCfg +from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import ( + BELT_CENTER_X, + BELT_CENTER_Y, + BELT_HALF_STRAIGHT, + BELT_TOP_Z, + BELT_TURN_RADIUS, +) + + +def test_racetrack_and_sorting_tasks_share_the_policy_contract() -> None: + """Sorting extends the four-cube task without changing its scene or policy configuration.""" + task = gym.spec("IsaacContrib-Conveyor-Franka-Newton-Play-v0") + base_task = gym.spec("IsaacContrib-Conveyor-Franka-Newton-v0") + base_cfg = ConveyorFrankaEnvCfg() + base_scene = base_cfg.scene.to_dict() + cfg = ConveyorFrankaA09A12EnvCfg() + cfg.scene._configure_route_assets(cfg.commands.transfer.parcel_colors) + + assert task.kwargs["env_cfg_entry_point"].endswith(":ConveyorFrankaA09A12EnvCfg") + assert base_task.kwargs["env_cfg_entry_point"].endswith(":ConveyorFrankaEnvCfg") + assert task.kwargs["rsl_rl_cfg_entry_point"] == base_task.kwargs["rsl_rl_cfg_entry_point"] + assert base_cfg.scene.to_dict() == base_scene + assert base_cfg.scene.to_dict() == ConveyorFrankaEnvCfg().scene.to_dict() + assert [name for name in vars(base_cfg.scene) if name.startswith("cube_")] == [ + "cube_0", + "cube_1", + "cube_2", + "cube_3", + ] + assert base_cfg.conveyor_force.transported_body_count_per_env == 4 + assert base_cfg.commands.transfer.class_type.__name__ == "ConveyorTransferCommand" + assert cfg.commands.transfer.class_type.__name__ == "ConveyorSortCommand" + assert cfg.conveyor_force.transported_body_count_per_env == 24 + assert cfg.scene.num_envs == 1 + assert cfg.actions == base_cfg.actions + assert cfg.observations == base_cfg.observations + assert cfg.commands.transfer.parcel_destinations == (0, 1) * 12 + assert cfg.commands.transfer.parcel_colors == ("blue", "orange", "green", "purple") * 6 + assert cfg.commands.transfer.randomize_arrivals is True + assert cfg.commands.transfer.hold_steps == base_cfg.commands.transfer.hold_steps + assert cfg.events == base_cfg.events + assert cfg.rewards == base_cfg.rewards + for name, term in vars(base_cfg.terminations).items(): + if name != "cube_out_of_workspace": + assert getattr(cfg.terminations, name) == term + assert cfg.terminations.cube_out_of_workspace.func == base_cfg.terminations.cube_out_of_workspace.func + assert cfg.decimation == base_cfg.decimation + assert cfg.sim.dt == base_cfg.sim.dt + assert cfg.sim.physics.load_visual_shapes is False + + +def test_sorting_contacts_exclude_stationary_conveyor_pairs() -> None: + """Touching belt sections must not consume contacts needed to support parcels.""" + import newton + import warp as wp + + builder = newton.ModelBuilder() + for x in (0.0, 0.15): + section = newton.ModelBuilder() + section.add_shape_box(-1, hx=0.1, hy=0.1, hz=0.1) + builder.add_builder(section, xform=wp.transform((x, 0.0, 0.0), wp.quat_identity())) + parcel = builder.add_body(xform=wp.transform((0.0, 0.0, 0.18), wp.quat_identity())) + builder.add_shape_box(parcel, hx=0.1, hy=0.1, hz=0.1) + model = builder.finalize(device="cpu") + cfg = ConveyorFrankaA09A12EnvCfg().sim.physics.collision_cfg + pipeline = newton.CollisionPipeline(model, **cfg.to_pipeline_args()) + contacts = pipeline.contacts() + pipeline.collide(model.state(), contacts) + count = int(contacts.rigid_contact_count.numpy()[0]) + assert count > 0 + bodies = model.shape_body.numpy() + first = bodies[contacts.rigid_contact_shape0.numpy()[:count]] + second = bodies[contacts.rigid_contact_shape1.numpy()[:count]] + assert ((first == parcel) | (second == parcel)).all() + + +def test_a09_a12_config_import_does_not_preload_usd() -> None: + """Task discovery must not import USD before a requested Kit application starts.""" + code = ( + "import sys; " + "import isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_asset_env_cfg; " + "raise SystemExit('pxr' in sys.modules)" + ) + result = subprocess.run([sys.executable, "-c", code], check=False) + assert result.returncode == 0 + + +def test_visual_only_usd_strips_physics_and_execution_metadata() -> None: + """Decorative references cannot introduce bodies, contacts, joints, or action graphs.""" + from pxr import Usd, UsdPhysics + + stage = Usd.Stage.CreateInMemory() + root = stage.DefinePrim("/VisualAsset", "Xform") + body = stage.DefinePrim("/VisualAsset/Body", "Xform") + shape = stage.DefinePrim("/VisualAsset/Body/Shape", "Cube") + joint = UsdPhysics.FixedJoint.Define(stage, "/VisualAsset/Joint").GetPrim() + graph = stage.DefinePrim("/VisualAsset/ActionGraph", "OmniGraph") + physics_scene = UsdPhysics.Scene.Define(stage, "/VisualAsset/PhysicsScene").GetPrim() + + UsdPhysics.ArticulationRootAPI.Apply(root) + UsdPhysics.RigidBodyAPI.Apply(body) + UsdPhysics.MassAPI.Apply(body) + UsdPhysics.CollisionAPI.Apply(shape) + UsdPhysics.MeshCollisionAPI.Apply(shape) + UsdPhysics.FilteredPairsAPI.Apply(shape) + shape.AddAppliedSchema("PhysxCollisionAPI") + + _make_usd_subtree_visual_only(root) + + assert not root.HasAPI(UsdPhysics.ArticulationRootAPI) + assert not body.HasAPI(UsdPhysics.RigidBodyAPI) + assert not body.HasAPI(UsdPhysics.MassAPI) + assert not shape.HasAPI(UsdPhysics.CollisionAPI) + assert not shape.HasAPI(UsdPhysics.MeshCollisionAPI) + assert not shape.HasAPI(UsdPhysics.FilteredPairsAPI) + assert "PhysxCollisionAPI" not in shape.GetAppliedSchemas() + assert not joint.IsActive() + assert not graph.IsActive() + assert not physics_scene.IsActive() + + +def test_digital_twin_assets_use_separate_physical_routes() -> None: + """Imported visual assets do not own contacts on the extended physical routes.""" + scene = ConveyorFrankaA09A12EnvCfg().scene + + assert scene.conveyor_left_belt_visual is None + assert scene.guard_left_inner_visual is None + assert hasattr(scene, "conveyor_left_top_straight_collision") + assert hasattr(scene, "conveyor_left_right_turn_collision") + assert hasattr(scene, "guard_left_inner_collision") + + asset_names = tuple( + name + for name in vars(scene) + if name.endswith(("_a09_visual", "_a12_visual")) and getattr(scene, name) is not None + ) + assert len(asset_names) == 2 + assert all(name.endswith("_a09_visual") for name in asset_names) + assert len(scene.build_conveyor_belt_specs()) > 8 + + for name in asset_names: + asset = getattr(scene, name) + assert asset.spawn.usd_path.endswith("conveyor_straight_supported.usd") + assert asset.spawn.collision_props is None + assert asset.spawn.make_uninstanceable + + +def test_thor_table_and_conveyors_rest_on_their_authored_supports() -> None: + """The Thor mount reaches the floor and narrow supports retain the conveyor deck elevation.""" + cfg = ConveyorFrankaA09A12EnvCfg() + scene = cfg.scene + scene._configure_route_assets() + + assert scene.tabletop.prim_path.endswith("/RobotThorTableVisual") + assert scene.tabletop.spawn.usd_path.endswith("/Props/Mounts/thor_table.usd") + assert scene.tabletop.spawn.collision_props is None + ground_z = 0.0 + assert scene.table_pedestal is None + assert math.isclose(scene.tabletop.init_state.pos[2] - 0.795 * scene.tabletop.spawn.scale[2], ground_z) + + for name in vars(scene): + if name.endswith(("_a09_visual", "_a12_visual")) and getattr(scene, name) is not None: + assert math.isclose(getattr(scene, name).init_state.pos[2], 0.55) + + assert ground_z == 0.0 + workspace_z = scene.ground.workspace_origin_offset[2] + assert 0.75 < workspace_z < 0.85 + assert math.isclose(scene.robot.init_state.pos[2], workspace_z) + assert math.isclose(scene.cube_0.init_state.pos[2], 0.28 + workspace_z) + base_collision_z = ConveyorFrankaEnvCfg().scene.conveyor_left_top_straight_collision.init_state.pos[2] + assert math.isclose(scene.warehouse_left_section_0.init_state.pos[2], base_collision_z + workspace_z) + + +def test_warehouse_layout_is_usd_authored_and_preserves_cube_physics() -> None: + """The presentation adds one USD assembly and changes only the task cubes' render spawner.""" + from pxr import Sdf + + scene = ConveyorFrankaA09A12EnvCfg().scene + base = ConveyorFrankaEnvCfg().scene + assert scene.warehouse_visual.spawn.collision_props is None + layer = Sdf.Layer.FindOrOpen(scene.warehouse_visual.spawn.usd_path) + assert layer.defaultPrim == "Warehouse" + assert layer.GetPrimAtPath("/Warehouse/Lights/WorkcellSoftbox") + assert layer.GetPrimAtPath("/Warehouse/Transport/Divert") + assert layer.GetPrimAtPath("/Warehouse/Parcels/Parcel00") + assert not layer.GetPrimAtPath("/Warehouse/Cell/Riser0") + assert not layer.GetPrimAtPath("/Warehouse/Supports") + assert not layer.GetPrimAtPath("/Warehouse/NetworkParcels") + assert layer.GetPrimAtPath("/Warehouse/Scanner").attributes["xformOp:rotateZ"].default == 90 + assert all( + path.startswith("https://") + or path + in {"conveyor_straight_supported.usd", "conveyor_quarter_supported.usd", "conveyor_routes.usda", "parcel.usda"} + for path in layer.GetExternalReferences() + ) + for cube_id in range(4): + visual_spawn = getattr(scene, f"cube_{cube_id}").spawn.to_dict() + base_spawn = getattr(base, f"cube_{cube_id}").spawn.to_dict() + visual_spawn.pop("func") + visual_spawn.pop("parcel_usd_path") + base_spawn.pop("func") + assert visual_spawn == base_spawn + scene._configure_route_assets() + for cube_id in range(4, 24): + cube = getattr(scene, f"cube_{cube_id}") + actual = cube.spawn.to_dict() + expected = scene.cube_0.spawn.to_dict() + assert actual.pop("parcel_usd_path").endswith( + f"parcel_{ConveyorFrankaA09A12EnvCfg().commands.transfer.parcel_colors[cube_id]}.usda" + ) + expected.pop("parcel_usd_path") + assert actual == expected + assert cube.prim_path.endswith(f"/Cube{cube_id}") + + +@pytest.mark.parametrize( + "parcel_asset", + [ + "parcel.usda", + "parcel_blue.usda", + "parcel_orange.usda", + "parcel_green.usda", + "parcel_purple.usda", + ], +) +def test_carton_visual_is_centered_on_the_original_40_mm_collider(tmp_path, monkeypatch, parcel_asset) -> None: + """Asset normalization changes appearance without moving or resizing the grasp surface.""" + from pxr import Usd, UsdGeom, UsdPhysics + + from isaaclab_tasks.contrib.conveyor_franka import conveyor_franka_asset_env_cfg as asset_cfg + + # Measured unscaled Cardbox_A1 bounds, used as an offline stand-in for the remote mesh. + lower = (-0.34971755743026733, -0.260576993227005, 0.0) + upper = (0.34971755743026733, 0.2605747878551483, 0.5099270939826965) + source = Usd.Stage.CreateNew(str(tmp_path / "carton.usda")) + root = UsdGeom.Xform.Define(source, "/Carton").GetPrim() + source.SetDefaultPrim(root) + UsdPhysics.RigidBodyAPI.Apply(root) + mesh = UsdGeom.Cube.Define(source, "/Carton/Mesh") + mesh.CreateSizeAttr(1.0) + mesh.AddTranslateOp().Set(tuple((a + b) / 2 for a, b in zip(lower, upper))) + mesh.AddScaleOp().Set(tuple(b - a for a, b in zip(lower, upper))) + UsdPhysics.CollisionAPI.Apply(mesh.GetPrim()) + source.GetRootLayer().Save() + monkeypatch.setattr(asset_cfg, "retrieve_file_path", lambda path: str(tmp_path / "carton.usda")) + _presentation_layer.cache_clear() + try: + stage = sim_utils.create_new_stage() + cfg = ConveyorFrankaA09A12EnvCfg().scene.cube_0.spawn + cfg.parcel_usd_path = str(_PRESENTATION_ASSETS / parcel_asset) + prim = cfg.func("/Cube", cfg) + visual = stage.GetPrimAtPath("/Cube/CartonVisual") + bounds = UsdGeom.BBoxCache(Usd.TimeCode.Default(), ["default", "render"]).ComputeWorldBound(visual) + assert tuple(bounds.ComputeAlignedRange().GetMin()) == pytest.approx((-0.02,) * 3, abs=1e-8) + assert tuple(bounds.ComputeAlignedRange().GetMax()) == pytest.approx((0.02,) * 3, abs=1e-8) + assert prim.HasAPI(UsdPhysics.RigidBodyAPI) + collider = stage.GetPrimAtPath("/Cube/geometry/mesh") + assert collider.HasAPI(UsdPhysics.CollisionAPI) + assert UsdGeom.Imageable(collider).ComputeVisibility() == "invisible" + assert sum(p.HasAPI(UsdPhysics.RigidBodyAPI) for p in stage.Traverse()) == 1 + assert sum(p.HasAPI(UsdPhysics.CollisionAPI) for p in stage.Traverse()) == 1 + assert not any(p.HasAPI(UsdPhysics.RigidBodyAPI) for p in Usd.PrimRange(visual)) + # The authored reference remains portable; cache resolution edits only an anonymous copy. + from pxr import Sdf + + assert all( + path.startswith("https://") + for path in Sdf.Layer.FindOrOpen(str(_PRESENTATION_ASSETS / "parcel.usda")).GetExternalReferences() + ) + finally: + _presentation_layer.cache_clear() + + +def test_asset_transforms_preserve_the_original_workcell() -> None: + """The working straights and adjoining quarter-turns keep their original locations and radii.""" + from pxr import Sdf + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import belt_collision_section_specs + from isaaclab_tasks.contrib.conveyor_franka.conveyor_warehouse_geometry import ( + warehouse_belt_sections, + ) + + scene = ConveyorFrankaA09A12EnvCfg().scene + layer = Sdf.Layer.FindOrOpen(scene.warehouse_visual.spawn.usd_path) + for side, sign, original_index in (("Left", 1, 1), ("Right", -1, 0)): + sections = warehouse_belt_sections(side, velocity=0.35) + original = belt_collision_section_specs(side)[original_index].geometry + inner = sections[0].geometry + assert inner.position == original.position + assert inner.size == original.size + visual = getattr(scene, f"conveyor_{side.lower()}_{'bottom' if sign > 0 else 'top'}_a09_visual") + assert visual.init_state.pos[:2] == ( + BELT_CENTER_X + BELT_HALF_STRAIGHT, + sign * (BELT_CENTER_Y - BELT_TURN_RADIUS), + ) + assert 4 * visual.spawn.scale[0] == pytest.approx(2 * BELT_HALF_STRAIGHT) + assert visual.init_state.pos[2] + 1.78053 * visual.spawn.scale[2] == pytest.approx( + BELT_TOP_Z + scene.ground.workspace_origin_offset[2] + ) + for name, x in ( + ("RightInner", BELT_CENTER_X + BELT_HALF_STRAIGHT), + ("LeftInner", BELT_CENTER_X - BELT_HALF_STRAIGHT), + ): + turn = next(section.belt for section in sections if f"{name}Turn" in section.geometry.name) + assert turn.pivot_point == (x, sign * BELT_CENTER_Y, 0) + assert turn.radius == BELT_TURN_RADIUS + authored = layer.GetPrimAtPath(f"/Warehouse/WorkcellConveyors/{side}{name}") + position = authored.attributes["xformOp:translate"].default + scale = authored.attributes["xformOp:scale"].default + angle = math.radians(authored.attributes["xformOp:rotateZ"].default) + # A03's original circular pivot is at (0, -1.4961); the referenced arc lands on the physical bend. + radius = 1.4961 * scale[0] + assert radius == pytest.approx(BELT_TURN_RADIUS) + assert position[0] + radius * math.sin(angle) == pytest.approx(x) + assert position[1] - radius * math.cos(angle) == pytest.approx(sign * BELT_CENTER_Y) + assert any(section.belt.direction[2] > 0 for section in sections if not section.belt.curved) + assert any(section.belt.direction[2] < 0 for section in sections if not section.belt.curved) + feed_velocity = 0.043 if side == "Left" else 0.052 + assert all( + section.belt.velocity == (feed_velocity if "Supply" in section.geometry.name else 0.35) + for section in sections + ) + + +@pytest.mark.parametrize("use_slice", [False, True]) +def test_warehouse_reset_loads_mixed_feeds_only_in_selected_environments(monkeypatch, use_slice) -> None: + """A seeded reset shuffles all physical arrivals without disturbing other environments.""" + import torch + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_cube_pool import ConveyorCubePool + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env import ConveyorFrankaEnv + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_warehouse_env import ConveyorFrankaWarehouseEnv + from isaaclab_tasks.contrib.conveyor_franka.conveyor_warehouse_geometry import warehouse_parcel_positions + + origins = torch.tensor([[0.0, 0.0, 0.8], [10.0, 12.0, 0.8]]) + cubes = {} + before = [] + velocities = [] + positions = warehouse_parcel_positions() + for i in range(len(positions)): + pose = torch.tensor([[0.52, 0.27, 0.06, 0.0, 0.0, 0.0, 1.0]]).repeat(2, 1) + pose[:, :3] += origins + before.append(pose.clone()) + velocity = torch.ones(2, 6) + velocities.append(velocity) + + def write(root_pose, env_ids, state=pose): + state[env_ids] = root_pose + + def write_velocity(root_velocity, env_ids, state=velocity): + state[env_ids] = root_velocity + + cubes[f"cube_{i}"] = SimpleNamespace( + data=SimpleNamespace(root_pose_w=SimpleNamespace(torch=pose)), + write_root_pose_to_sim_index=write, + write_root_velocity_to_sim_index=write_velocity, + ) + + class Scene(dict): + env_origins = origins + _ALL_INDICES = torch.arange(2) + + env = ConveyorFrankaWarehouseEnv.__new__(ConveyorFrankaWarehouseEnv) + env.scene = Scene(cubes) + env.conveyor_cube_pool = ConveyorCubePool(tuple(cubes.values()), 2, "cpu") + env._warehouse_animation = [] + env.sim = SimpleNamespace(device="cpu", remove_render_callback=lambda name: None) + env.cfg = ConveyorFrankaA09A12EnvCfg() + monkeypatch.setattr(ConveyorFrankaEnv, "_reset_idx", lambda self, ids: None) + monkeypatch.setattr(ConveyorFrankaEnv, "close", lambda self: None) + torch.manual_seed(42) + env._reset_idx(slice(1, 2) if use_slice else torch.tensor([1])) + actual_positions = [] + for i in range(len(positions)): + actual = cubes[f"cube_{i}"].data.root_pose_w.torch + torch.testing.assert_close(actual[0], before[i][0]) + torch.testing.assert_close(actual[1, 3:], before[i][1, 3:]) + actual_positions.append(actual[1, :3] - origins[1]) + torch.testing.assert_close(velocities[i][0], torch.ones(6)) + torch.testing.assert_close(velocities[i][1], torch.zeros(6)) + actual_positions = torch.stack(actual_positions) + expected = torch.tensor(positions) + distances = torch.linalg.vector_norm(actual_positions[:, None] - expected[None], dim=-1) + assert distances.min(dim=1).values.max() < 1e-5 + assert distances.argmin(dim=1).unique().numel() == len(positions) + assert not torch.allclose(actual_positions, expected) + torch.manual_seed(42) + env._reset_idx(slice(1, 2) if use_slice else torch.tensor([1])) + repeated = torch.stack([cube.data.root_pose_w.torch[1, :3] - origins[1] for cube in cubes.values()]) + torch.testing.assert_close(repeated, actual_positions) + + +def test_warehouse_animation_uses_active_kit_viewer_and_policy_time(tmp_path, monkeypatch) -> None: + """A CLI-selected Kit viewer animates and loops authored parcels without changing task state.""" + from pxr import Usd, UsdGeom + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env import ConveyorFrankaEnv + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_warehouse_env import ConveyorFrankaWarehouseEnv + + path = str(tmp_path / "traffic.usda") + source = Usd.Stage.CreateNew(path) + root = UsdGeom.Xform.Define(source, "/Warehouse").GetPrim() + source.SetDefaultPrim(root) + source.SetTimeCodesPerSecond(60) + source.SetEndTimeCode(120) + parcel = UsdGeom.Xform.Define(source, "/Warehouse/Parcels/Parcel00") + translate = parcel.AddTranslateOp() + translate.Set((0, 0, 1), 0) + translate.Set((2, 0, 1), 120) + parcel.AddRotateXYZOp().Set((0, 0, 0)) + source.GetRootLayer().Save() + stage = Usd.Stage.CreateInMemory() + stage.DefinePrim("/World/envs/env_0/WarehouseVisual").GetReferences().AddReference(path) + callbacks = {} + + def initialize(env, cfg, **kwargs): + env.cfg = cfg + env.common_step_counter = 60 + env.scene = SimpleNamespace(env_prim_paths=["/World/envs/env_0"]) + env.sim = SimpleNamespace( + stage=stage, + visualizers=[SimpleNamespace(cfg=SimpleNamespace(visualizer_type="kit"))], + set_setting=lambda name, value: None, + add_render_callback=lambda name, callback: callbacks.update({name: callback}), + remove_render_callback=lambda name: callbacks.pop(name, None), + ) + + monkeypatch.setattr(ConveyorFrankaEnv, "__init__", initialize) + monkeypatch.setattr(ConveyorFrankaEnv, "close", lambda env: None) + cfg = ConveyorFrankaA09A12EnvCfg() + # CLI viewer selection is resolved by SimulationContext, not written back into this list. + cfg.sim.visualizer_cfgs = [] + cfg.scene.warehouse_visual.spawn.usd_path = path + env = ConveyorFrankaWarehouseEnv(cfg) + try: + callback = callbacks["conveyor_warehouse_animation"] + callback(None) + position = stage.GetPrimAtPath("/World/envs/env_0/WarehouseVisual/Parcels/Parcel00").GetAttribute( + "xformOp:translate" + ) + assert tuple(position.Get()) == pytest.approx((1, 0, 1)) + env.common_step_counter = 180 + callback(None) + assert tuple(position.Get()) == pytest.approx((1, 0, 1)) + finally: + env.close() + _presentation_layer.cache_clear() + assert not callbacks + + +@pytest.mark.parametrize("remote_position", [(2.8, 0.1, 0.56), (1.2, 0.59, 0.16)]) +def test_warehouse_policy_view_preserves_local_states_and_physical_inventory(remote_position): + """Only remote transport is mapped to waiting slots; physical tensors remain untouched.""" + import torch + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_warehouse_env import ConveyorFrankaWarehouseEnv + + origin = torch.tensor([[10.0, 12.0, 0.8]]) + positions = ( + torch.tensor([[[0.52, 0.27, 0.06], [0.6, -0.1, 0.25], remote_position, [1.9, -1.1, 0.2]]]) + origin[:, None] + ) + quaternions = torch.randn(1, 4, 4) + velocities = torch.randn(1, 4, 6) + original = tuple(value.clone() for value in (positions, quaternions, velocities)) + env = SimpleNamespace( + scene=SimpleNamespace(env_origins=origin), + device="cpu", + cfg=SimpleNamespace(conveyor_force=SimpleNamespace(speed=0.35)), + _in_workcell=ConveyorFrankaWarehouseEnv._in_workcell, + ) + actual = ConveyorFrankaWarehouseEnv._adapt_policy_cube_state(env, positions, quaternions, velocities) + for result, source, before in zip(actual, (positions, quaternions, velocities), original): + torch.testing.assert_close(result[:, :2], source[:, :2]) + torch.testing.assert_close(source, before) + torch.testing.assert_close(actual[0][0, 2:, 1:] - origin[0, 1:], torch.tensor([[0.75, 0.06], [-0.75, 0.06]])) + torch.testing.assert_close(actual[1][0, 2:], torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]])) + + +def test_warehouse_idle_parks_and_preserves_invalid_actions(monkeypatch): + """Idle handling cannot mask invalid policy input or bypass an active transfer.""" + import torch + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env import ConveyorFrankaEnv + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_warehouse_env import ConveyorFrankaWarehouseEnv + + class Scene(dict): + num_envs = 3 + + scene = Scene( + robot=SimpleNamespace( + data=SimpleNamespace( + default_joint_pos=SimpleNamespace(torch=torch.zeros(3, 7)), + joint_pos=SimpleNamespace(torch=torch.full((3, 7), 0.12)), + ) + ) + ) + command = SimpleNamespace(has_target=torch.tensor([False, True, False])) + env = ConveyorFrankaWarehouseEnv.__new__(ConveyorFrankaWarehouseEnv) + env.scene = scene + env.sim = SimpleNamespace(device="cpu", remove_render_callback=lambda name: None) + env.cfg = SimpleNamespace(actions=SimpleNamespace(arm_action=SimpleNamespace(scale=0.12))) + env.command_manager = SimpleNamespace(get_term=lambda name: command) + env._warehouse_arm_joint_ids = list(range(7)) + env._warehouse_animation = [] + accepted = [] + result = ({"policy": torch.zeros(3, 123)}, None, None, None, {}) + + def step(self, action): + accepted.append(action) + return result + + monkeypatch.setattr(ConveyorFrankaEnv, "step", step) + monkeypatch.setattr(ConveyorFrankaEnv, "close", lambda self: None) + action = torch.full((3, 8), 0.5) + action[2, 0] = torch.nan + assert env.step(action) is result + actual = accepted[0] + torch.testing.assert_close(actual[0], torch.tensor([-0.25] * 7 + [0.0])) + torch.testing.assert_close(actual[1], action[1]) + torch.testing.assert_close(actual[2], action[2], equal_nan=True) + + +def test_parcel_pool_preserves_pinned_slots_and_eventually_assigns_every_physical_parcel(): + """Arrivals rotate through four distinct slots without displacing a held parcel or another environment.""" + import torch + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_cube_pool import ConveyorCubePool + + pool = ConveyorCubePool((None,) * 16, 2, "cpu") + pool.slot_ids[1] = torch.tensor([3, 2, 1, 0]) + positions = torch.zeros(2, 16, 3) + positions[:, :, 0] = torch.linspace(0.4, 1.0, 16) + local = torch.zeros(2, 16, dtype=torch.bool) + local[0, 1] = True + candidates = torch.zeros_like(local) + candidates[0, 4:] = True + target_slots = torch.tensor([0, 2]) + pinned = torch.tensor([True, True]) + initial_other_environment = pool.slot_ids[1].clone() + for _ in range(6): + changed = pool.refresh(positions, local, candidates, target_slots, pinned) + assert changed.tolist() == [True, False] + assert pool.slot_ids[0, :2].tolist() == [0, 1] + assert len(pool.slot_ids[0].unique()) == 4 + torch.testing.assert_close(pool.slot_ids[1], initial_other_environment) + assert bool((pool.assignment_counts[0] > 0).all()) + physical_id = int(pool.slot_ids[0, 2]) + pool.record_transfers(torch.tensor([0]), torch.tensor([2])) + assert pool.transfer_counts[0, physical_id] == 1 + assert pool.transfer_counts.sum() == 1 + pool.reset(torch.tensor([0])) + assert pool.slot_ids[0].tolist() == [0, 1, 2, 3] + torch.testing.assert_close(pool.slot_ids[1], initial_other_environment) + + +def test_parcel_pool_gathers_per_environment_states_and_checks_unassigned_inventory(): + """All policy features use the same physical assignment; an unassigned fallen cube still terminates.""" + import torch + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_cube_pool import ConveyorCubePool, cube_values + from isaaclab_tasks.contrib.conveyor_franka.mdp.observations import _cube_state + from isaaclab_tasks.contrib.conveyor_franka.mdp.terminations import cube_out_of_workspace + + assets = [] + for cube_id in range(6): + values = { + "root_pos_w": torch.tensor([[0.1 * cube_id, 0.27, 0.06], [10 + 0.1 * cube_id, 0.27, 0.06]]), + "root_quat_w": torch.full((2, 4), float(cube_id)), + "root_vel_w": torch.full((2, 6), float(cube_id)), + } + assets.append( + SimpleNamespace( + data=SimpleNamespace(**{name: SimpleNamespace(torch=value) for name, value in values.items()}) + ) + ) + pool = ConveyorCubePool(tuple(assets), 2, "cpu") + pool.slot_ids[:] = torch.tensor([[4, 1, 2, 3], [0, 5, 2, 3]]) + env = SimpleNamespace( + conveyor_cube_pool=pool, scene=SimpleNamespace(env_origins=torch.tensor([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]])) + ) + physical_before = cube_values(env, "root_pos_w", all_cubes=True).clone() + states = _cube_state(env) + for values, attribute in zip(states, ("root_pos_w", "root_quat_w", "root_vel_w")): + for row in range(2): + expected = torch.stack([getattr(assets[i].data, attribute).torch[row] for i in pool.slot_ids[row]]) + torch.testing.assert_close(values[row], expected) + torch.testing.assert_close(cube_values(env, "root_pos_w", all_cubes=True), physical_before) + assert not cube_out_of_workspace(env).any() + assets[5].data.root_pos_w.torch[0, 2] = -1.0 # Parcel 5 is not assigned in environment 0. + assert cube_out_of_workspace(env).tolist() == [True, False] diff --git a/source/isaaclab_tasks/test/contrib/test_conveyor_franka_geometry.py b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_geometry.py new file mode 100644 index 000000000000..c49409122aa3 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_geometry.py @@ -0,0 +1,120 @@ +# 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 + +"""Durable geometry checks for the contributed conveyor Franka task.""" + +from collections import Counter + +import pytest + +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env_cfg import _collision_properties, _cube +from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import ( + BELT_TURN_RADIUS, + TURN_SEGMENT_COUNT, + MeshSpec, + belt_collision_section_specs, + belt_mesh_spec, + guard_mesh_specs, +) + + +def _edge_use_counts(spec: MeshSpec) -> Counter[tuple[int, int]]: + """Count triangle uses of every undirected mesh edge.""" + edges: Counter[tuple[int, int]] = Counter() + for triangle in spec.faces: + for start, end in zip(triangle, triangle[1:] + triangle[:1], strict=True): + edges[tuple(sorted((start, end)))] += 1 + return edges + + +def test_racetrack_visual_meshes_are_named_watertight_loops(): + """Belts and rails remain uniquely named, closed racetrack meshes.""" + specs = tuple(spec for side in ("Left", "Right") for spec in (belt_mesh_spec(side), *guard_mesh_specs(side))) + expected_loop_vertices = 2 * TURN_SEGMENT_COUNT + 2 + + assert len(specs) == 6 + assert len({spec.name for spec in specs}) == len(specs) + for spec in specs: + assert len(spec.vertices) == 4 * expected_loop_vertices + assert len(spec.faces) == 8 * expected_loop_vertices + assert set(_edge_use_counts(spec).values()) == {2} + + +@pytest.mark.parametrize("warehouse", [False, True]) +def test_belt_top_faces_point_upward(warehouse): + """One-sided triangle-mesh surfaces support parcels from above.""" + for side in ("Left", "Right"): + if warehouse: + from isaaclab_tasks.contrib.conveyor_franka.conveyor_warehouse_geometry import ( + warehouse_belt_sections, + warehouse_guard_meshes, + ) + + specs = [ + section.geometry for section in warehouse_belt_sections(side) if isinstance(section.geometry, MeshSpec) + ] + specs.extend(warehouse_guard_meshes(side)) + else: + specs = [belt_mesh_spec(side)] + for spec in specs: + assert set(_edge_use_counts(spec).values()) == {2} + _assert_top_faces_point_upward(spec) + + +def _assert_top_faces_point_upward(spec): + top_z = max(vertex[2] for vertex in spec.vertices) + for face in spec.faces: + vertices = tuple(spec.vertices[index] for index in face) + if all(vertex[2] == top_z for vertex in vertices): + a, b, c = vertices + cross_z = (b[0] - a[0]) * (c[1] - a[1]) - (b[1] - a[1]) * (c[0] - a[0]) + assert cross_z > 0.0 + + +def test_contact_configuration_uses_one_mujoco_parameterization(): + """Raw MuJoCo solref must not be combined with shadowed Newton force-space gains.""" + mujoco_cfg = _collision_properties()[-1] + cube_material = _cube("TestCube", (1.0, 0.0, 0.0), (0.0, 0.0, 0.0)).spawn.physics_material + + assert mujoco_cfg.solref is not None + assert cube_material.contact_stiffness is None + assert cube_material.contact_damping is None + assert cube_material.torsional_friction is None + assert cube_material.rolling_friction is None + + +def test_collision_sections_carry_schema_aligned_belt_intent(): + """Task geometry and runtime descriptions share paths, units, and curve semantics.""" + sections = belt_collision_section_specs("Left", velocity=0.35, friction_coefficient=0.5, contact_threshold=0.997) + + assert len(sections) == 4 + assert tuple(section.belt.prim_path for section in sections) == tuple( + f"{{ENV_REGEX_NS}}/{section.geometry.name}" for section in sections + ) + assert tuple(section.belt.velocity for section in sections) == (0.35,) * 4 + assert tuple(section.belt.friction_coefficient for section in sections) == (0.5,) * 4 + assert tuple(section.belt.contact_threshold for section in sections) == (0.997,) * 4 + assert tuple(section.belt.curved for section in sections) == (False, False, True, True) + assert tuple(section.belt.radius for section in sections) == (None, None, BELT_TURN_RADIUS, BELT_TURN_RADIUS) + + +def test_elevated_belt_normals_accept_contacts_on_every_ramp_panel(): + """Inclined parcels receive traction instead of sliding into the bottom transition.""" + import numpy as np + + from isaaclab_tasks.contrib.conveyor_franka.conveyor_warehouse_geometry import warehouse_belt_sections + + for side in ("Left", "Right"): + for section in warehouse_belt_sections(side): + if section.belt.curved or abs(section.belt.direction[2]) < 1e-6: + continue + vertices = np.asarray(section.geometry.vertices) + triangles = vertices[np.asarray(section.geometry.faces)] + normals = np.cross(triangles[:, 1] - triangles[:, 0], triangles[:, 2] - triangles[:, 0]) + top = normals[normals[:, 2] > 1e-8] + top /= np.linalg.norm(top, axis=1, keepdims=True) + assert len(top) > 0 + assert np.all(top @ section.belt.surface_normal >= section.belt.contact_threshold) + assert abs(np.dot(section.belt.direction, section.belt.surface_normal)) < 1e-6 diff --git a/source/isaaclab_tasks/test/contrib/test_conveyor_franka_mdp.py b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_mdp.py new file mode 100644 index 000000000000..9a8d50b2d4b1 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_mdp.py @@ -0,0 +1,648 @@ +# 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 + +"""Unit tests for conveyor-transfer state, curriculum, and success geometry.""" + +from collections import Counter +from types import SimpleNamespace + +import pytest +import torch + +from isaaclab_tasks.contrib.conveyor_franka.agents.rsl_rl_ppo_cfg import ( + ConveyorFrankaPPORunnerCfg, + ConveyorGaussianBernoulliDistribution, +) +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env_cfg import ConveyorFrankaEnvCfg +from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import ( + BELT_CENTER_X, + BELT_INNER_STRAIGHT_Y, + BELT_OUTER_STRAIGHT_Y, +) +from isaaclab_tasks.contrib.conveyor_franka.mdp.actions import ( + ConveyorRelativeJointPositionAction, + ResetBufferedGripperAction, +) +from isaaclab_tasks.contrib.conveyor_franka.mdp.commands import ConveyorTransferCommand, transfer_success_mask +from isaaclab_tasks.contrib.conveyor_franka.mdp.curriculums import ( + deployment_probability_from_progress, + reset_sampling_probabilities, +) +from isaaclab_tasks.contrib.conveyor_franka.mdp.observations import classify_cube_conveyors +from isaaclab_tasks.contrib.conveyor_franka.mdp.reset_events import ( + CUBE_COUNT, + ConveyorResetRecipe, + _balanced_cube_slots, + _sample_collision_free_active_x, + build_reset_rows, + franka_tool_position, + reset_variant_counts, + select_next_transfer_cube, +) +from isaaclab_tasks.contrib.conveyor_franka.mdp.rewards import finite_action_rate_l2, transfer_potential +from isaaclab_tasks.contrib.conveyor_franka.mdp.terminations import ( + invalid_action, + subgoal_time_out, + transfer_sequence_time_out, +) + + +def _make_arm_action_term() -> ConveyorRelativeJointPositionAction: + """Build the tensor-only portion of the arm action term without a simulator.""" + workspace_lower = torch.tensor((-0.75, -0.45, -0.55, -2.75, -0.45, 1.85, -0.10)) + workspace_upper = torch.tensor((0.85, 0.85, 0.35, -1.75, 0.45, 3.05, 1.65)) + positions = ((workspace_lower + workspace_upper) * 0.5).repeat(2, 1) + limits = torch.tensor((-4.0, 4.0)).repeat(2, 7, 1) + action = object.__new__(ConveyorRelativeJointPositionAction) + action.cfg = SimpleNamespace(max_delta=0.12, joint_limit_margin=0.02, clip=None) + action._asset = SimpleNamespace( + data=SimpleNamespace( + joint_pos=SimpleNamespace(torch=positions), + soft_joint_pos_limits=SimpleNamespace(torch=limits), + ) + ) + action._joint_ids = slice(None) + action._scale = 0.12 + action._offset = 0.0 + action._workspace_lower = workspace_lower + action._workspace_upper = workspace_upper + action._raw_actions = torch.zeros((2, 7)) + action._previous_actions = torch.zeros((2, 7)) + action._processed_actions = torch.zeros((2, 7)) + action._position_targets = positions.clone() + action._invalid_actions = torch.zeros(2, dtype=torch.bool) + return action + + +def _make_gripper_action_term() -> ResetBufferedGripperAction: + """Build the tensor-only portion of the binary gripper action term.""" + action = object.__new__(ResetBufferedGripperAction) + action.cfg = SimpleNamespace(clip=None, command_name="transfer", force_close_steps=2) + action._raw_actions = torch.zeros((2, 1)) + action._previous_actions = torch.zeros((2, 1)) + action._processed_actions = torch.zeros((2, 2)) + action._open_command = torch.full((2,), 0.04) + action._close_command = torch.zeros(2) + action._invalid_actions = torch.zeros(2, dtype=torch.bool) + command = SimpleNamespace(held_cube_ids=torch.full((2,), -1, dtype=torch.long)) + action._env = SimpleNamespace( + command_manager=SimpleNamespace(get_term=lambda _name: command), + episode_length_buf=torch.zeros(2, dtype=torch.long), + ) + return action + + +def test_arm_action_sanitizes_nonfinite_values_and_clamps_normalized_input(): + """Invalid policy outputs cannot reach joint targets or exceed one normalized unit.""" + action = _make_arm_action_term() + policy_actions = torch.tensor(((float("nan"), float("inf"), -float("inf"), 5.0, -5.0, 0.5, -0.5), (-2.0,) * 7)) + + action.process_actions(policy_actions) + + expected = torch.tensor(((0.0, 1.0, -1.0, 1.0, -1.0, 0.5, -0.5), (-1.0,) * 7)) + torch.testing.assert_close(action.raw_actions, expected) + assert action.invalid_actions.tolist() == [True, False] + assert torch.isfinite(action.processed_actions).all() + expected_targets = action._asset.data.joint_pos.torch + expected * 0.12 + torch.testing.assert_close(action.processed_actions, expected_targets) + + +def test_gripper_action_sanitizes_nonfinite_values_before_binary_mapping(): + """The eighth policy dimension cannot silently turn a NaN into an open command.""" + action = _make_gripper_action_term() + + action.process_actions(torch.tensor(((float("nan"),), (float("inf"),)))) + + torch.testing.assert_close(action.raw_actions, torch.tensor(((-1.0,), (1.0,)))) + assert action.invalid_actions.tolist() == [True, True] + torch.testing.assert_close(action.processed_actions[0], action._close_command) + torch.testing.assert_close(action.processed_actions[1], action._open_command) + + +def test_invalid_action_termination_and_reset_are_per_environment(): + """Arm and gripper failures aggregate per environment, and reset clears the arm flag.""" + arm_action = _make_arm_action_term() + gripper_action = _make_gripper_action_term() + arm_action._invalid_actions[:] = torch.tensor((True, False)) + gripper_action._invalid_actions[:] = torch.tensor((False, True)) + actions = {"arm_action": arm_action, "gripper_action": gripper_action} + env = SimpleNamespace(action_manager=SimpleNamespace(get_term=actions.__getitem__)) + + assert invalid_action(env).tolist() == [True, True] + arm_action.reset([0]) + gripper_action._invalid_actions[1] = False + + assert invalid_action(env).tolist() == [False, False] + + +def test_action_rate_uses_finite_commands_and_preserves_invalid_termination(): + """A rejected policy output has a finite final reward without hiding its termination.""" + arm_action = _make_arm_action_term() + gripper_action = _make_gripper_action_term() + arm_action.process_actions(torch.full((2, 7), 0.25)) + gripper_action.process_actions(torch.tensor(((-1.0,), (1.0,)))) + arm_action.process_actions( + torch.tensor(((float("nan"), float("inf"), -float("inf"), 5.0, -5.0, 0.5, -0.5), (0.5,) * 7)) + ) + gripper_action.process_actions(torch.tensor(((float("nan"),), (-1.0,)))) + actions = {"arm_action": arm_action, "gripper_action": gripper_action} + env = SimpleNamespace(num_envs=2, action_manager=SimpleNamespace(get_term=actions.__getitem__)) + + reward = finite_action_rate_l2(env) + + expected_arm = torch.square(arm_action.raw_actions - arm_action.previous_actions).sum(dim=1) + expected_gripper = torch.square(gripper_action.raw_actions - gripper_action.previous_actions).sum(dim=1) + torch.testing.assert_close(reward, expected_arm + expected_gripper) + assert torch.isfinite(reward).all() + assert invalid_action(env).tolist() == [True, False] + + +def test_action_rate_matches_standard_l2_for_ordinary_policy_actions(): + """Finite in-range policy commands retain the standard action-rate semantics.""" + arm_action = _make_arm_action_term() + gripper_action = _make_gripper_action_term() + previous_arm = torch.tensor(((0.1,) * 7, (-0.2,) * 7)) + current_arm = torch.tensor(((-0.3,) * 7, (0.4,) * 7)) + previous_gripper = torch.tensor(((-1.0,), (1.0,))) + current_gripper = -previous_gripper + arm_action.process_actions(previous_arm) + gripper_action.process_actions(previous_gripper) + arm_action.process_actions(current_arm) + gripper_action.process_actions(current_gripper) + actions = {"arm_action": arm_action, "gripper_action": gripper_action} + env = SimpleNamespace(num_envs=2, action_manager=SimpleNamespace(get_term=actions.__getitem__)) + + reward = finite_action_rate_l2(env) + + previous = torch.cat((previous_arm, previous_gripper), dim=1) + current = torch.cat((current_arm, current_gripper), dim=1) + torch.testing.assert_close(reward, torch.square(current - previous).sum(dim=1)) + + +def test_final_config_validation_catches_overridden_arm_contracts(): + """Top-level validation runs after overrides and protects workspace-to-joint alignment.""" + cfg = ConveyorFrankaEnvCfg() + assert cfg.seed is None + cfg.validate() + + cfg.actions.arm_action.preserve_order = False + try: + cfg.validate() + except ValueError as exc: + assert "preserve" in str(exc) + else: + raise AssertionError("Expected invalid arm ordering to fail configuration validation.") + + +def test_production_solver_and_viewer_defaults_are_bounded(): + """The task keeps CUDA graphs enabled and avoids scene-wide over-allocation or rendering.""" + cfg = ConveyorFrankaEnvCfg() + physics = cfg.sim.physics + solver = physics.solver_cfg + + assert cfg.conveyor_force.speed == 0.35 + assert physics.use_cuda_graph is True + assert physics.load_visual_shapes is None + assert solver.njmax == 300 + assert solver.nconmax == 200 + assert solver.impratio == 1.0 + assert cfg.sim.default_visualizer_cfg.max_visible_envs == 1 + assert cfg.sim.default_visualizer_cfg.randomly_sample_visible_envs is False + + +def test_reset_rows_cover_every_cube_direction_and_phase_once(): + """The reset bank is the complete command and physical-phase cross product.""" + rows = build_reset_rows() + + assert len(rows) == sum(reset_variant_counts()) * CUBE_COUNT * 2 + assert Counter((row.recipe, row.variant_id, row.target_cube_id, row.source_side_id) for row in rows) == Counter( + (recipe, variant_id, cube_id, side_id) + for recipe in ConveyorResetRecipe + for variant_id in range(reset_variant_counts()[int(recipe)]) + for cube_id in range(CUBE_COUNT) + for side_id in range(2) + ) + for row in rows: + expected_held = row.recipe in { + ConveyorResetRecipe.LIFT, + ConveyorResetRecipe.CARRY, + ConveyorResetRecipe.PLACE, + } or ( + row.recipe == ConveyorResetRecipe.GRASP + and row.variant_id == reset_variant_counts()[int(ConveyorResetRecipe.GRASP)] - 1 + ) + assert row.held == expected_held + + grasp_rows = [ + row + for row in rows + if row.recipe == ConveyorResetRecipe.GRASP and row.target_cube_id == 0 and row.source_side_id == 0 + ] + assert all(first.finger_position > second.finger_position for first, second in zip(grasp_rows, grasp_rows[1:])) + assert grasp_rows[-1].held + + +def test_reset_arm_anchors_reach_expected_transfer_waypoints(): + """IK anchors place the tool over source, transit, or destination waypoints.""" + anchor_variants = { + ConveyorResetRecipe.GOAL: 0, + ConveyorResetRecipe.PLACE: 3, + ConveyorResetRecipe.CARRY: 2, + ConveyorResetRecipe.LIFT: 3, + ConveyorResetRecipe.GRASP: 0, + ConveyorResetRecipe.PREGRASP: 0, + } + rows = [ + row + for row in build_reset_rows() + if row.target_cube_id == 0 and anchor_variants.get(row.recipe) == row.variant_id + ] + joints = torch.tensor([row.arm_positions for row in rows], dtype=torch.float64) + positions = franka_tool_position(joints) + + for row, position in zip(rows, positions, strict=True): + source_y = 0.27 if row.source_side_id == 0 else -0.27 + target_y = -source_y + if row.recipe == ConveyorResetRecipe.BELT: + continue + expected = { + ConveyorResetRecipe.PREGRASP: (0.52, source_y, 0.14), + ConveyorResetRecipe.GRASP: (0.52, source_y, 0.06), + ConveyorResetRecipe.LIFT: (0.52, source_y, 0.22), + ConveyorResetRecipe.CARRY: (0.52, 0.0, 0.25), + ConveyorResetRecipe.PLACE: (0.52, target_y, 0.105), + ConveyorResetRecipe.GOAL: (0.52, target_y, 0.14), + }[row.recipe] + torch.testing.assert_close(position, torch.tensor(expected, dtype=position.dtype), atol=3.0e-4, rtol=0.0) + + +def test_cube_conveyor_state_has_stable_three_way_encoding(): + """Cube side observations distinguish both belts from the transfer corridor.""" + positions = torch.tensor([[[0.5, 0.27, 0.06], [0.5, 0.0, 0.20], [0.5, -0.27, 0.06]]]) + + encoded = classify_cube_conveyors(positions) + + torch.testing.assert_close( + encoded, + torch.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]]), + ) + + +def test_released_goal_state_succeeds_only_on_commanded_conveyor(): + """The same physical cube placement is successful for exactly one direction.""" + cube_positions = torch.tensor([[0.58, -0.27, 0.06], [0.58, -0.27, 0.06]]) + cube_velocities = torch.zeros((2, 3)) + tool_positions = torch.tensor([[0.58, -0.27, 0.14], [0.58, -0.27, 0.14]]) + finger_positions = torch.full((2, 2), 0.04) + target_side_ids = torch.tensor([1, 0]) + + successful = transfer_success_mask( + cube_positions, + cube_velocities, + tool_positions, + finger_positions, + target_side_ids, + ) + + assert successful.tolist() == [True, False] + + +def test_transfer_potential_increases_through_release(): + """Dense shaping must not punish lowering and releasing at the destination.""" + source_side = torch.zeros(5, dtype=torch.long) + cube_positions = torch.tensor( + [ + [0.52, 0.27, 0.06], + [0.52, 0.27, 0.20], + [0.52, 0.00, 0.20], + [0.52, -0.27, 0.105], + [0.52, -0.27, 0.06], + ] + ) + tool_positions = cube_positions.clone() + tool_positions[-1, 2] += 0.10 + finger_positions = torch.full((5, 2), 0.019) + finger_positions[-1] = 0.04 + + potentials = transfer_potential(cube_positions, tool_positions, finger_positions, source_side) + + assert torch.all(potentials[1:] > potentials[:-1]) + + +def test_transfer_potential_rewards_closing_only_near_cube(): + """The acquisition bridge credits a close command only around the object.""" + cube_positions = torch.tensor([[0.52, 0.27, 0.06], [0.52, 0.27, 0.06]]) + tool_positions = torch.tensor([[0.52, 0.27, 0.07], [0.52, 0.27, 0.20]]) + source_side = torch.zeros(2, dtype=torch.long) + open_fingers = torch.full((2, 2), 0.04) + closed_fingers = torch.full((2, 2), 0.019) + + open_potential = transfer_potential(cube_positions, tool_positions, open_fingers, source_side) + closed_potential = transfer_potential(cube_positions, tool_positions, closed_fingers, source_side) + + assert closed_potential[0] - open_potential[0] > 0.5 + assert closed_potential[1] - open_potential[1] < 0.03 + + +def test_policy_distribution_samples_exact_binary_gripper_and_finite_kl(): + """PPO likelihoods match the gripper command that reaches physics.""" + distribution = ConveyorGaussianBernoulliDistribution(output_dim=8) + distribution.update(torch.zeros((4096, 8))) + + samples = distribution.sample() + old_params = tuple(parameter.clone() for parameter in distribution.params) + distribution.update(torch.full((4096, 8), 0.2)) + divergence = distribution.kl_divergence(old_params, distribution.params) + + assert set(torch.unique(samples[:, -1]).tolist()) == {-1.0, 1.0} + assert torch.isfinite(distribution.log_prob(samples)).all() + assert torch.isfinite(divergence).all() + assert torch.all(divergence >= 0.0) + + +def test_reset_sampling_guarantees_deployment_mass_and_tracks_frontier(): + """Sampling reserves deployment starts and favors intermediate frontier rows.""" + rows = build_reset_rows() + recipe_ids = torch.tensor([row.recipe for row in rows], dtype=torch.long) + variant_ids = torch.tensor([row.variant_id for row in rows], dtype=torch.long) + target_cube_ids = torch.tensor([row.target_cube_id for row in rows], dtype=torch.long) + source_side_ids = torch.tensor([row.source_side_id for row in rows], dtype=torch.long) + place_stratum = (recipe_ids == int(ConveyorResetRecipe.PLACE)) & (target_cube_ids == 0) & (source_side_ids == 0) + place_ids = torch.nonzero(place_stratum, as_tuple=False).flatten() + monitor_cfg = ConveyorFrankaEnvCfg().curriculum.reset_sampling.params["success_monitor"] + monitor = monitor_cfg.class_type(monitor_cfg, num_partitions=1, partition_size=len(rows), device="cpu") + monitor.success_update( + torch.cat((place_ids[0].repeat(50), place_ids[1].repeat(50))), + torch.cat((torch.ones(50, dtype=torch.bool), torch.arange(50) % 2 == 0)), + ) + + deployment_rows = (recipe_ids == int(ConveyorResetRecipe.BELT)) & ( + variant_ids == reset_variant_counts()[int(ConveyorResetRecipe.BELT)] - 1 + ) + probabilities = reset_sampling_probabilities( + recipe_ids, + variant_ids, + target_cube_ids, + source_side_ids, + monitor.target_weights(), + deployment_probability=0.35, + ) + + torch.testing.assert_close(probabilities.sum(), torch.tensor(1.0)) + torch.testing.assert_close(probabilities[deployment_rows].sum(), torch.tensor(0.35)) + torch.testing.assert_close(probabilities[~deployment_rows].sum(), torch.tensor(0.65)) + assert probabilities[place_ids[0]] < probabilities[place_ids[1]] + for recipe in ConveyorResetRecipe: + for cube_id in range(CUBE_COUNT): + for side_id in range(2): + stratum_rows = (recipe_ids == int(recipe)) & (target_cube_ids == cube_id) & (source_side_ids == side_id) + torch.testing.assert_close( + probabilities[stratum_rows & ~deployment_rows].sum(), + torch.tensor(0.65 / (len(ConveyorResetRecipe) * CUBE_COUNT * 2)), + ) + torch.testing.assert_close( + probabilities[deployment_rows & (source_side_ids == 0)].sum(), + torch.tensor(0.35 / 2), + ) + torch.testing.assert_close( + probabilities[deployment_rows & (source_side_ids == 1)].sum(), + torch.tensor(0.35 / 2), + ) + + +def test_deployment_probability_increases_with_rolling_readiness(): + """Mastered, well-covered reset rows shift sampling toward deployment starts.""" + kwargs = { + "initial_probability": 0.35, + "final_probability": 0.90, + "progress_start": 0.45, + "progress_end": 0.80, + "coverage_target": 0.50, + } + + initial = deployment_probability_from_progress(torch.tensor(0.30), torch.tensor(1.0), **kwargs) + middle = deployment_probability_from_progress(torch.tensor(0.625), torch.tensor(0.50), **kwargs) + final = deployment_probability_from_progress(torch.tensor(0.90), torch.tensor(1.0), **kwargs) + + torch.testing.assert_close(initial, torch.tensor(0.35)) + torch.testing.assert_close(middle, torch.tensor(0.625)) + torch.testing.assert_close(final, torch.tensor(0.90)) + + +def test_next_transfer_cube_is_random_among_eligible_alternatives(): + """Continuing commands use the one-hot target instead of a cyclic identity shortcut.""" + count = 4096 + positions = torch.zeros((count, CUBE_COUNT, 3)) + positions[:, :, 1] = torch.tensor((-0.27, -0.27, -0.27, 0.27)) + current_cube_ids = torch.zeros(count, dtype=torch.long) + source_side_ids = torch.ones(count, dtype=torch.long) + torch.manual_seed(7) + + selected = select_next_transfer_cube(positions, current_cube_ids, source_side_ids) + + assert set(selected.tolist()) == {1, 2} + frequencies = torch.bincount(selected, minlength=CUBE_COUNT).float() / count + assert torch.all(torch.abs(frequencies[1:3] - 0.5) < 0.05) + + +@pytest.mark.parametrize("use_pool", [False, True]) +def test_manual_transfer_goal_uses_selected_cube_current_side(use_pool): + """Viewer goal changes preserve cube identity and infer the opposite destination.""" + + class _Scene(dict): + pass + + origins = torch.tensor(((0.0, 1.0, 0.0), (0.0, -2.0, 0.0))) + selected_cube_positions = torch.tensor(((0.2, 1.3, 0.06), (0.7, -2.3, 0.06))) + cubes = tuple( + SimpleNamespace( + data=SimpleNamespace( + root_pos_w=SimpleNamespace( + torch=selected_cube_positions if cube_id == 2 else torch.zeros_like(selected_cube_positions) + ) + ) + ) + for cube_id in range(CUBE_COUNT) + ) + scene = _Scene() + scene.env_origins = origins + command = object.__new__(ConveyorTransferCommand) + command._env = SimpleNamespace( + num_envs=2, + device="cpu", + scene=scene, + episode_length_buf=torch.tensor((11, 19)), + ) + for cube_id, cube in enumerate(cubes): + scene[f"cube_{cube_id}"] = cube + if use_pool: + from isaaclab_tasks.contrib.conveyor_franka.conveyor_cube_pool import ConveyorCubePool + + extra_cube = SimpleNamespace(data=SimpleNamespace(root_pos_w=SimpleNamespace(torch=selected_cube_positions))) + cubes[2].data.root_pos_w.torch = 2 * origins - selected_cube_positions + pool = ConveyorCubePool((*cubes, extra_cube), 2, "cpu") + pool.slot_ids[:, 2] = 4 + command._env.conveyor_cube_pool = pool + command.target_cube_ids = torch.tensor((0, 1)) + command.source_side_ids = torch.tensor((1, 0)) + command.held_cube_ids = torch.tensor((0, 1)) + command.subgoal_start_steps = torch.tensor((2, 3)) + command._stable_steps = torch.ones(2, dtype=torch.long) + command.is_success = torch.ones(2, dtype=torch.bool) + command.new_success = torch.ones(2, dtype=torch.bool) + command.pending_success = torch.ones(2, dtype=torch.bool) + command._last_evaluation_steps = torch.tensor((11, 19)) + + command.set_goal(2) + + assert command.target_cube_ids.tolist() == [2, 2] + assert command.source_side_ids.tolist() == [0, 1] + assert command.held_cube_ids.tolist() == [-1, -1] + assert command.subgoal_start_steps.tolist() == [11, 19] + torch.testing.assert_close( + command.command, + torch.tensor(((0, 0, 1, 0, 0, 1), (0, 0, 1, 0, 1, 0)), dtype=torch.float32), + ) + + +def test_continuing_training_truncates_only_stalled_or_long_sequences(): + """Subgoal and sequence limits bound training without ending successful transfers.""" + assert not ConveyorFrankaPPORunnerCfg().init_at_random_ep_len + command = SimpleNamespace( + evaluate=lambda: None, + subgoal_start_steps=torch.tensor((0, 0, 300)), + pending_success=torch.zeros(3, dtype=torch.bool), + transfer_counts=torch.tensor((0, 7, 8)), + ) + env = SimpleNamespace( + step_dt=0.1, + episode_length_buf=torch.tensor((199, 200, 450)), + command_manager=SimpleNamespace(get_term=lambda _name: command), + ) + + assert subgoal_time_out(env, timeout_s=20.0).tolist() == [False, True, False] + assert transfer_sequence_time_out(env, maximum_transfers=8).tolist() == [False, False, True] + + +def test_active_cube_sampling_avoids_inactive_source_lane_cubes(): + """Random deployment starts cannot begin with interpenetrating parcels.""" + count = 2048 + base_x = torch.tensor((0.26, 0.42, 0.72, 0.88)).expand(count, -1).clone() + cube_sides = torch.tensor((0, 0, 1, 1)).expand(count, -1).clone() + target_cube_ids = torch.arange(count) % CUBE_COUNT + source_side_ids = torch.arange(count) % 2 + cube_sides.scatter_(1, target_cube_ids.unsqueeze(1), source_side_ids.unsqueeze(1)) + + sampled = _sample_collision_free_active_x( + base_x, + cube_sides, + target_cube_ids, + source_side_ids, + torch.full((count,), 0.30), + torch.full((count,), 0.82), + ) + + cube_ids = torch.arange(CUBE_COUNT).expand(count, -1) + inactive_on_source = (cube_sides == source_side_ids.unsqueeze(1)) & (cube_ids != target_cube_ids.unsqueeze(1)) + separation = torch.abs(sampled.unsqueeze(1) - base_x) + assert torch.all(separation[inactive_on_source] >= 0.055) + + +def test_deployment_layout_uses_each_racetrack_straight_run_once(): + """Every command keeps its target reachable and distributes all four cubes evenly.""" + target_cube_ids = torch.arange(CUBE_COUNT).repeat_interleave(2) + source_side_ids = torch.arange(2).repeat(CUBE_COUNT) + + slots, cube_sides, cube_x, cube_y = _balanced_cube_slots(target_cube_ids, source_side_ids) + + expected_slots = torch.arange(CUBE_COUNT).expand(target_cube_ids.numel(), -1) + torch.testing.assert_close(torch.sort(slots, dim=1).values, expected_slots) + active_slots = slots.gather(1, target_cube_ids.unsqueeze(1)).squeeze(1) + torch.testing.assert_close(active_slots, 2 * source_side_ids) + assert torch.all(torch.sum(cube_sides == 0, dim=1) == 2) + assert torch.all(torch.sum(cube_sides == 1, dim=1) == 2) + + expected_y = torch.tensor( + (-BELT_OUTER_STRAIGHT_Y, -BELT_INNER_STRAIGHT_Y, BELT_INNER_STRAIGHT_Y, BELT_OUTER_STRAIGHT_Y) + ) + torch.testing.assert_close(torch.sort(cube_y, dim=1).values, expected_y.expand_as(cube_y)) + for side_id in (0, 1): + side_x = torch.where(cube_sides == side_id, cube_x, torch.nan) + expected_sum = torch.full((source_side_ids.numel(),), 2 * BELT_CENTER_X, dtype=cube_x.dtype) + torch.testing.assert_close(torch.nansum(side_x, dim=1), expected_sum) + + +def test_sort_dispatch_preserves_grasp_and_never_reverses_a_sorted_parcel(monkeypatch): + """Dispatch changes logical slots, never physical state, and leaves completed batches circulating.""" + from isaaclab_tasks.contrib.conveyor_franka.conveyor_cube_pool import ConveyorCubePool + from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_warehouse_env import ConveyorFrankaWarehouseEnv + from isaaclab_tasks.contrib.conveyor_franka.mdp import sorting + + count = 6 + positions = torch.tensor([[[2.0, 0.8, 0.46]] * count] * 2) + # The wrong-class arrival is outside the initial four slots; the nearby natural carton is already sorted. + positions[0, 0] = torch.tensor([0.7, 0.27, 0.06]) + positions[0, 5] = torch.tensor([0.8, 0.27, 0.06]) + positions[1, 1] = torch.tensor([0.5, -0.1, 0.25]) # Active grasp crossing between loops. + assets = tuple( + SimpleNamespace( + data=SimpleNamespace( + root_pos_w=SimpleNamespace(torch=positions[:, i]), + root_lin_vel_w=SimpleNamespace(torch=torch.zeros(2, 3)), + ) + ) + for i in range(count) + ) + pool = ConveyorCubePool(assets, 2, "cpu") + env = SimpleNamespace( + num_envs=2, + device="cpu", + scene=SimpleNamespace(env_origins=torch.zeros(2, 3)), + conveyor_cube_pool=pool, + _in_workcell=ConveyorFrankaWarehouseEnv._in_workcell, + episode_length_buf=torch.tensor([20, 20]), + ) + command = object.__new__(sorting.ConveyorSortCommand) + command._env = env + command.cfg = sorting.ConveyorSortCommandCfg(parcel_destinations=(0, 1) * 3) + command.parcel_destinations = torch.tensor(command.cfg.parcel_destinations) + command.has_target = torch.tensor([False, True]) + command.target_cube_ids = torch.tensor([0, 1]) + command.source_side_ids = torch.tensor([0, 0]) + command.held_cube_ids = torch.full((2,), -1) + command.subgoal_start_steps = torch.zeros(2, dtype=torch.long) + command.command_counter = torch.ones(2, dtype=torch.long) + command._stable_steps = torch.zeros(2, dtype=torch.long) + command._last_evaluation_steps = torch.full((2,), -1) + command.is_success = torch.zeros(2, dtype=torch.bool) + command.new_success = torch.zeros(2, dtype=torch.bool) + command.pending_success = torch.zeros(2, dtype=torch.bool) + command.metrics = {name: torch.zeros(2) for name in ("sorted_parcels", "batch_complete")} + monkeypatch.setattr(sorting, "physical_cube_acquisition_mask", lambda *args, **kwargs: torch.tensor([False, True])) + before = positions.clone() + command._update_command() + assert command.has_target.tolist() == [True, True] + assert pool.slot_ids[0, command.target_cube_ids[0]] == 5 + assert pool.slot_ids[1, command.target_cube_ids[1]] == 1 + assert command.command[0, -2:].tolist() == [0.0, 1.0] + assert command.metrics["sorted_parcels"].tolist() == [1.0, 0.0] + torch.testing.assert_close(positions, before) + # Every parcel is now on its class's loop. Stable completion must end dispatch, not reverse direction. + positions[0, :, 0] = 0.6 + positions[0, :, 1] = torch.tensor([0.27, -0.27] * 3) + positions[0, :, 2] = 0.06 + command.pending_success[0] = True + command._update_command() + assert command.has_target.tolist() == [False, True] + assert command.metrics["batch_complete"].tolist() == [1.0, 0.0] + command._update_command() + assert not command.has_target[0] + # An idle command must not keep paying the preceding placement reward. + command.has_target.zero_() + command.new_success.fill_(True) + command.is_success.fill_(True) + command.evaluate() + assert not command.new_success.any() + assert not command.is_success.any() diff --git a/source/isaaclab_tasks/test/contrib/test_conveyor_franka_physx_cfg.py b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_physx_cfg.py new file mode 100644 index 000000000000..10bdb1e105b6 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_conveyor_franka_physx_cfg.py @@ -0,0 +1,107 @@ +# 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 + +"""Import-light checks for the checkpoint-compatible PhysX conveyor configuration.""" + +import gymnasium as gym +import pytest +from isaaclab_physx.physics import PhysxCfg +from isaaclab_physx.sim.spawners.materials import PhysxRigidBodyMaterialCfg + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_env_cfg import ConveyorFrankaEnvCfg +from isaaclab_tasks.contrib.conveyor_franka.conveyor_franka_physx_env_cfg import ( + ConveyorFrankaPhysxEnvCfg, + physx_belt_section_specs, +) +from isaaclab_tasks.contrib.conveyor_franka.conveyor_geometry import BELT_TURN_RADIUS, MeshSpec +from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + + +def test_physx_task_is_registered_with_a_dedicated_config() -> None: + """The native backend is opt-in and cannot alter the Newton task registration.""" + physx_spec = gym.spec("IsaacContrib-Conveyor-Franka-PhysX-CPU-v0") + newton_spec = gym.spec("IsaacContrib-Conveyor-Franka-Newton-v0") + + assert physx_spec.entry_point == newton_spec.entry_point + assert physx_spec.kwargs["env_cfg_entry_point"].endswith(":ConveyorFrankaPhysxEnvCfg") + assert newton_spec.kwargs["env_cfg_entry_point"].endswith(":ConveyorFrankaEnvCfg") + + +def test_physx_config_preserves_policy_and_timing_contracts() -> None: + """A Newton checkpoint sees the same ordered 8-D action and 60 Hz policy interface.""" + newton_cfg = ConveyorFrankaEnvCfg() + physx_cfg = ConveyorFrankaPhysxEnvCfg() + + newton_cfg.validate() + physx_cfg.validate() + assert isinstance(physx_cfg.sim.physics, PhysxCfg) + assert physx_cfg.sim.device == "cpu" + assert physx_cfg.scene.num_envs == 1 + assert physx_cfg.sim.dt == newton_cfg.sim.dt == 1.0 / 120.0 + assert physx_cfg.decimation == newton_cfg.decimation == 2 + assert physx_cfg.actions.arm_action.joint_names == newton_cfg.actions.arm_action.joint_names + assert physx_cfg.actions.gripper_action.joint_names == newton_cfg.actions.gripper_action.joint_names + assert physx_cfg.conveyor_force.speed == newton_cfg.conveyor_force.speed == 0.35 + assert physx_cfg.scene.robot.spawn.joint_drive_props is None + assert physx_cfg.scene.robot.spawn.rigid_props.disable_gravity is True + + +def test_physx_task_default_device_survives_an_unset_parser_override() -> None: + """Default-backend callers can preserve the task's declared CPU device.""" + cfg = parse_env_cfg("IsaacContrib-Conveyor-Franka-PhysX-CPU-v0", device=None, num_envs=2) + + assert cfg.sim.device == "cpu" + assert cfg.scene.num_envs == 2 + + +@pytest.mark.parametrize("device", ["cuda", "cuda:0", "cuda:1"]) +def test_physx_config_rejects_broken_gpu_surface_velocity_contacts(device: str) -> None: + """The pinned Isaac Sim GPU path must not silently let cubes tunnel through belts.""" + cfg = ConveyorFrankaPhysxEnvCfg() + cfg.sim.device = device + + with pytest.raises(ValueError, match="CPU-only.*--device cpu"): + cfg.validate() + + +def test_curved_physx_sections_are_pivot_local_watertight_sdfs() -> None: + """Turn roots sit at their pivots so native angular velocity has the intended center.""" + cfg = ConveyorFrankaPhysxEnvCfg() + sections = physx_belt_section_specs("Left", velocity=0.35) + + assert len(sections) == 4 + for section, root_position in sections[2:]: + assert isinstance(section.geometry, MeshSpec) + assert section.belt.curved + assert section.belt.radius == BELT_TURN_RADIUS + assert section.belt.pivot_point == (0.0, 0.0, 0.0) + assert root_position[2] == 0.0 + + turn_asset = cfg.scene.conveyor_left_right_turn_collision + assert turn_asset.spawn.collision_approximation == "sdf" + assert isinstance(turn_asset.spawn.physics_material, PhysxRigidBodyMaterialCfg) + assert turn_asset.spawn.physics_material.dynamic_friction == 0.5 + assert turn_asset.spawn.rigid_props.kinematic_enabled is True + + +def test_runtime_specs_and_material_override_have_one_source_of_truth() -> None: + """Final overrides reach both the runtime view and native PhysX material authoring.""" + cfg = ConveyorFrankaPhysxEnvCfg() + cfg.scene.configure_conveyor(friction_coefficient=0.42) + specs = cfg.scene.build_conveyor_belt_specs( + velocity=0.21, + friction_coefficient=0.42, + contact_threshold=0.95, + ) + + assert len(specs) == 8 + assert {spec.velocity for spec in specs} == {0.21} + assert {spec.friction_coefficient for spec in specs} == {0.42} + assert {spec.contact_threshold for spec in specs} == {0.95} + for side in ("left", "right"): + for section_key in ("top_straight", "bottom_straight", "right_turn", "left_turn"): + asset = getattr(cfg.scene, f"conveyor_{side}_{section_key}_collision") + assert asset.spawn.physics_material.dynamic_friction == 0.42 diff --git a/source/isaaclab_tasks/test/env_test_utils.py b/source/isaaclab_tasks/test/env_test_utils.py index d72a998ffb8e..31dccf0e285e 100644 --- a/source/isaaclab_tasks/test/env_test_utils.py +++ b/source/isaaclab_tasks/test/env_test_utils.py @@ -231,7 +231,7 @@ def _configure_osc_smoke_actions(env, actions: torch.Tensor) -> None: def _run_environments( task_name, - device, + device: str | None, num_envs, num_steps=20, multi_agent=False, @@ -243,7 +243,7 @@ def _run_environments( Args: task_name: Name of the environment. - device: Device to use (e.g., 'cuda'). + device: Device override to use (e.g., 'cuda'), or None to preserve the task default. num_envs: Number of environments. num_steps: Number of simulation steps. multi_agent: Whether the environment is multi-agent. @@ -284,7 +284,7 @@ def _run_environments( def _check_random_actions( task_name: str, - device: str, + device: str | None, num_envs: int, num_steps: int = 20, multi_agent: bool = False, @@ -296,7 +296,7 @@ def _check_random_actions( Args: task_name: Name of the environment. - device: Device to use (e.g., 'cuda'). + device: Device override to use (e.g., 'cuda'), or None to preserve the task default. num_envs: Number of environments. num_steps: Number of simulation steps. multi_agent: Whether the environment is multi-agent. diff --git a/source/isaaclab_visualizers/changelog.d/maximiliank-newton-ui-callback.rst b/source/isaaclab_visualizers/changelog.d/maximiliank-newton-ui-callback.rst new file mode 100644 index 000000000000..4f2dd1181e2c --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/maximiliank-newton-ui-callback.rst @@ -0,0 +1,4 @@ +Added +^^^^^ + +* Exposed Newton viewer ImGui callback registration through the Isaac Lab visualizer. diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py index 0897e3d19869..ac915b9d622e 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py @@ -15,7 +15,7 @@ from collections.abc import Callable from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import numpy as np # noqa: F401 — used in type hints and colorization helpers import torch @@ -1390,6 +1390,24 @@ def supports_live_plots(self) -> bool: """Newton RTX viewers do not provide live-plot panels; GL viewers do.""" return False + def register_ui_callback( + self, + callback: Callable[[Any], None], + position: Literal["side", "stats", "free", "panel", "rendering"] = "side", + ) -> None: + """Register a callback with the active Newton viewer's ImGui interface. + + Args: + callback: Function called during UI rendering with the active ImGui module. + position: Newton viewer UI region in which the callback is rendered. + + Raises: + RuntimeError: If the visualizer has not created an interactive viewer. + """ + if self._viewer is None: + raise RuntimeError("Newton UI callbacks require an initialized interactive viewer.") + self._viewer.register_ui_callback(callback, position=position) + def is_training_paused(self) -> bool: """Return whether training is paused from viewer controls.""" if not self._is_initialized or self._viewer is None: diff --git a/source/isaaclab_visualizers/test/test_newton_adapter.py b/source/isaaclab_visualizers/test/test_newton_adapter.py index 4306acc33554..499e3bb52e4b 100644 --- a/source/isaaclab_visualizers/test/test_newton_adapter.py +++ b/source/isaaclab_visualizers/test/test_newton_adapter.py @@ -217,6 +217,17 @@ def __init__(self): assert visualizer.cfg.lookat == (0.0, 0.0, 1.0) +def test_newton_visualizer_register_ui_callback_forwards_to_active_viewer(): + callback = Mock() + viewer = SimpleNamespace(register_ui_callback=Mock()) + visualizer = NewtonGLVisualizer(NewtonGLVisualizerCfg()) + visualizer._viewer = viewer + + visualizer.register_ui_callback(callback, position="panel") + + viewer.register_ui_callback.assert_called_once_with(callback, position="panel") + + @pytest.mark.parametrize( "cfg_type", [NewtonGLVisualizerCfg, NewtonRTXVisualizerCfg, RerunVisualizerCfg, ViserVisualizerCfg] ) diff --git a/uv.lock b/uv.lock index 2ab99273c008..7303adba8f48 100644 --- a/uv.lock +++ b/uv.lock @@ -1698,7 +1698,7 @@ wheels = [ [[package]] name = "isaaclab" -version = "29.0.0" +version = "31.0.0" source = { editable = "source/isaaclab" } [[package]] @@ -1722,7 +1722,7 @@ requires-dist = [ [[package]] name = "isaaclab-contrib" -version = "2.0.1" +version = "3.0.0" source = { editable = "source/isaaclab_contrib" } dependencies = [ { name = "isaaclab" }, @@ -2066,7 +2066,7 @@ provides-extras = ["tetrahedralization", "video", "importers", "test", "dev", "s [[package]] name = "isaaclab-experimental" -version = "0.2.2" +version = "0.2.3" source = { editable = "source/isaaclab_experimental" } dependencies = [ { name = "isaaclab" }, @@ -2077,7 +2077,7 @@ requires-dist = [{ name = "isaaclab", editable = "source/isaaclab" }] [[package]] name = "isaaclab-mimic" -version = "2.0.9" +version = "2.0.11" source = { editable = "source/isaaclab_mimic" } dependencies = [ { name = "isaaclab" }, @@ -2094,7 +2094,7 @@ requires-dist = [ [[package]] name = "isaaclab-newton" -version = "8.0.0" +version = "9.0.0" source = { editable = "source/isaaclab_newton" } dependencies = [ { name = "isaaclab" }, @@ -2105,7 +2105,7 @@ requires-dist = [{ name = "isaaclab", editable = "source/isaaclab" }] [[package]] name = "isaaclab-ov" -version = "3.4.0" +version = "3.5.0" source = { editable = "source/isaaclab_ov" } dependencies = [ { name = "isaaclab" }, @@ -2120,7 +2120,7 @@ requires-dist = [ [[package]] name = "isaaclab-physx" -version = "7.2.3" +version = "7.3.0" source = { editable = "source/isaaclab_physx" } dependencies = [ { name = "isaaclab" }, @@ -2142,7 +2142,7 @@ requires-dist = [{ name = "isaaclab", editable = "source/isaaclab" }] [[package]] name = "isaaclab-rl" -version = "1.2.0" +version = "1.3.0" source = { editable = "source/isaaclab_rl" } dependencies = [ { name = "isaaclab" }, @@ -2159,17 +2159,21 @@ requires-dist = [ [[package]] name = "isaaclab-tasks" -version = "21.1.0" +version = "21.3.0" source = { editable = "source/isaaclab_tasks" } dependencies = [ { name = "isaaclab" }, { name = "isaaclab-assets" }, + { name = "isaaclab-newton" }, + { name = "isaaclab-physx" }, ] [package.metadata] requires-dist = [ { name = "isaaclab", editable = "source/isaaclab" }, { name = "isaaclab-assets", editable = "source/isaaclab_assets" }, + { name = "isaaclab-newton", editable = "source/isaaclab_newton" }, + { name = "isaaclab-physx", editable = "source/isaaclab_physx" }, ] [[package]] @@ -2200,7 +2204,7 @@ requires-dist = [{ name = "isaaclab", editable = "source/isaaclab" }] [[package]] name = "isaaclab-visualizers" -version = "1.13.0" +version = "2.0.0" source = { editable = "source/isaaclab_visualizers" } dependencies = [ { name = "isaaclab" },