diff --git a/docs/source/tutorial/index.rst b/docs/source/tutorial/index.rst index 5357ec04..7ccb59a2 100644 --- a/docs/source/tutorial/index.rst +++ b/docs/source/tutorial/index.rst @@ -13,23 +13,24 @@ Follow the tutorials in this order for the best learning experience: 1. :doc:`create_scene` — Set up a simulation, add objects, and run the render loop. **Start here.** 2. :doc:`create_softbody` and :doc:`create_cloth` — Add deformable bodies to your scenes. 3. :doc:`rigid_object_group` — Manage collections of rigid objects efficiently. -4. :doc:`robot` — Load and control a robot in simulation. -5. :doc:`sensor` — Add cameras and capture RGB/depth/segmentation data. -6. :doc:`solver` — Configure IK solvers for end-effector control. -7. :doc:`motion_gen` — Generate smooth trajectories with motion planners. -8. :doc:`atomic_actions` — Use built-in action primitives (move, move joints, pick, move held object, place). -9. :doc:`gizmo` — Interactively control robots with on-screen gizmos. +4. :doc:`rigid_constraint` — Attach and detach two rigid objects via a fixed constraint. +5. :doc:`robot` — Load and control a robot in simulation. +6. :doc:`sensor` — Add cameras and capture RGB/depth/segmentation data. +7. :doc:`solver` — Configure IK solvers for end-effector control. +8. :doc:`motion_gen` — Generate smooth trajectories with motion planners. +9. :doc:`atomic_actions` — Use built-in action primitives (move, move joints, pick, move held object, place). +10. :doc:`gizmo` — Interactively control robots with on-screen gizmos. **Phase 2: Environments** -10. :doc:`basic_env` — Create a simple Gymnasium environment with ``BaseEnv``. Prerequisite: Phase 1 basics. -11. :doc:`modular_env` — Build a config-driven environment with ``EmbodiedEnv``, managers, and randomization. Prerequisite: :doc:`basic_env`. -12. :doc:`data_generation` — Generate expert demonstration datasets for imitation learning. Prerequisite: :doc:`modular_env`. -13. :doc:`rl` — Train RL agents with PPO or GRPO. Prerequisite: :doc:`basic_env`. +11. :doc:`basic_env` — Create a simple Gymnasium environment with ``BaseEnv``. Prerequisite: Phase 1 basics. +12. :doc:`modular_env` — Build a config-driven environment with ``EmbodiedEnv``, managers, and randomization. Prerequisite: :doc:`basic_env`. +13. :doc:`data_generation` — Generate expert demonstration datasets for imitation learning. Prerequisite: :doc:`modular_env`. +14. :doc:`rl` — Train RL agents with PPO or GRPO. Prerequisite: :doc:`basic_env`. **Phase 3: Extending the Framework** -14. :doc:`add_robot` — Add a new robot model to EmbodiChain. +15. :doc:`add_robot` — Add a new robot model to EmbodiChain. .. toctree:: :maxdepth: 1 @@ -39,6 +40,7 @@ Follow the tutorials in this order for the best learning experience: create_softbody create_cloth rigid_object_group + rigid_constraint robot add_robot solver diff --git a/docs/source/tutorial/rigid_constraint.rst b/docs/source/tutorial/rigid_constraint.rst new file mode 100644 index 00000000..500558a6 --- /dev/null +++ b/docs/source/tutorial/rigid_constraint.rst @@ -0,0 +1,158 @@ +Rigid constraint tutorial +========================== + +.. currentmodule:: embodichain.lab.sim + +This tutorial shows how to attach two rigid objects via a fixed physics +constraint, observe the constraint holding their relative pose, and then remove +it. It follows the style used in the :doc:`rigid_object_group` tutorial and +references the example script located in +``scripts/tutorials/sim/create_rigid_constraint.py``. + +A *fixed constraint* (a weld) binds two dynamic bodies so that their relative +pose is held constant by the physics solver — they move as a single rigid +assembly until the constraint is removed. This is useful for grasping and +assembly tasks, where an object must be "held" to a gripper or two parts must +be joined temporarily. + +.. tip:: + + Constraints are created and removed through the + :class:`SimulationManager`, which owns one constraint handle per arena. The + same API is exposed as on-demand event functors (``create_rigid_constraint`` + / ``remove_rigid_constraint`` in ``embodichain.lab.gym.envs.managers.events``) + so a task environment can attach/detach mid-episode via + ``env.event_manager.apply(mode="attach", env_ids=...)``. + +The Code +~~~~~~~~ + +The tutorial corresponds to the ``create_rigid_constraint.py`` script in the +``scripts/tutorials/sim`` directory. + +.. dropdown:: Code for create_rigid_constraint.py + :icon: code + + .. literalinclude:: ../../../scripts/tutorials/sim/create_rigid_constraint.py + :language: python + :linenos: + + +The Code Explained +~~~~~~~~~~~~~~~~~~ + + +Adding two cubes +---------------- + +Two dynamic cubes are added with :meth:`SimulationManager.add_rigid_object`. +Each uses a :class:`CubeCfg` shape (a primitive cube, so no mesh asset file is +needed) and a :class:`RigidBodyAttributesCfg` for mass and friction. ``cube_a`` +is placed slightly higher than ``cube_b`` so that, once detached, the lower +cube lands first and the relative pose visibly changes. + +.. literalinclude:: ../../../scripts/tutorials/sim/create_rigid_constraint.py + :language: python + :start-at: cube_a = sim.add_rigid_object( + :end-at: print("[INFO]: Scene setup complete with two cubes (cube_a, cube_b).") + + +Attaching the cubes +------------------- + +The two cubes are welded with :meth:`SimulationManager.create_rigid_constraint`. +A :class:`RigidConstraintCfg` names the constraint and points at the two object +UIDs. ``local_frame_a`` / ``local_frame_b`` default to ``None``, so the +constraint welds the cubes at their *current* relative pose: ``local_frame_a`` +defaults to identity (object A's origin) and ``local_frame_b`` is computed from +the objects' current poses so that the offset is preserved rather than the two +origins being pulled together. Pass explicit ``(4, 4)`` matrices — or an +``(N, 4, 4)`` array for one frame per arena — to define a specific joint frame +instead. + +.. literalinclude:: ../../../scripts/tutorials/sim/create_rigid_constraint.py + :language: python + :start-at: constraint = sim.create_rigid_constraint( + :end-at: print("[INFO]: Created constraint 'cube_weld' between cube_a and cube_b.") + +While attached, the cubes' relative pose stays essentially constant across +physics steps because the solver enforces the constraint. (``constraint.get_relative_transform()`` +returns the constraint-frame transform, which is ~0 while the constraint is +satisfied; the tutorial instead prints the bodies' relative z, ``cube_b.z - +cube_a.z``, to make the held offset visible.) + + +Removing the constraint +----------------------- + +The constraint is removed by name with +:meth:`SimulationManager.remove_rigid_constraint`. After removal, +:meth:`SimulationManager.get_rigid_constraint` returns ``None``, and the two +cubes are independent again — their relative pose is no longer enforced and +will drift as they interact with gravity and the ground. + +.. literalinclude:: ../../../scripts/tutorials/sim/create_rigid_constraint.py + :language: python + :start-at: sim.remove_rigid_constraint("cube_weld") + :end-at: print("\n[INFO]: Removed constraint 'cube_weld'. cube_a and cube_b are now free.") + +.. attention:: + + ``remove_rigid_constraint`` accepts an ``env_ids`` argument, so in a + vectorized simulation you can detach a subset of arenas while leaving the + rest attached. Likewise, ``create_rigid_constraint`` accepts ``env_ids`` to + attach only specific arenas. + + +Using the constraint from a task environment +-------------------------------------------- + +Inside a Gym environment the same operations are triggered on demand through +event functors registered under custom modes. A task wires up the attach and +detach functors, then calls ``event_manager.apply`` when its own logic decides +(for example, when a gripper closes or opens): + +.. code-block:: python + + from embodichain.lab.gym.envs.managers.cfg import EventCfg, SceneEntityCfg + from embodichain.lab.gym.envs.managers.events import ( + create_rigid_constraint, + remove_rigid_constraint, + ) + from embodichain.utils import configclass + + @configclass + class MyTaskEventsCfg: + attach_objects: EventCfg = EventCfg( + func=create_rigid_constraint, + mode="attach", + params={ + "obj_a_cfg": SceneEntityCfg(uid="cube_a"), + "obj_b_cfg": SceneEntityCfg(uid="cube_b"), + "name": "cube_weld", + }, + ) + detach_objects: EventCfg = EventCfg( + func=remove_rigid_constraint, + mode="detach", + params={"name": "cube_weld"}, + ) + + # Triggered from the task's own step / reset logic: + self.event_manager.apply(mode="attach", env_ids=gripping_env_ids) + self.event_manager.apply(mode="detach", env_ids=released_env_ids) + + +Running the tutorial +~~~~~~~~~~~~~~~~~~~~ + +To run the script from the repository root: + +.. code-block:: bash + + python scripts/tutorials/sim/create_rigid_constraint.py + +You can pass flags such as ``--headless``, ``--num_envs ``, and +``--device `` to customize the run. With the default settings the +script prints the cubes' relative z-position every 20 steps, first while +attached (held constant) and then after removal (free to drift). diff --git a/embodichain/lab/gym/envs/managers/events.py b/embodichain/lab/gym/envs/managers/events.py index e99f1458..a531ac40 100644 --- a/embodichain/lab/gym/envs/managers/events.py +++ b/embodichain/lab/gym/envs/managers/events.py @@ -20,6 +20,8 @@ import os import random +import numpy as np + from copy import deepcopy from typing import TYPE_CHECKING, List, Tuple, Dict @@ -30,7 +32,7 @@ Articulation, Robot, ) -from embodichain.lab.sim.cfg import RigidObjectCfg, ArticulationCfg +from embodichain.lab.sim.cfg import RigidObjectCfg, ArticulationCfg, RigidConstraintCfg from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg from embodichain.lab.gym.envs.managers import Functor, FunctorCfg @@ -619,3 +621,70 @@ def set_detached_uids_for_env_reset( """ env.add_detached_uids_for_reset(uids=uids) + + +def create_rigid_constraint( + env: EmbodiedEnv, + env_ids: torch.Tensor | None, + obj_a_cfg: SceneEntityCfg, + obj_b_cfg: SceneEntityCfg, + name: str, + local_frame_a: np.ndarray | None = None, + local_frame_b: np.ndarray | None = None, +) -> None: + """Attach two rigid objects via a fixed constraint for the given env_ids. + + Registered under a custom event mode (e.g. ``"attach"``); the task triggers it + with ``env.event_manager.apply(mode="attach", env_ids=...)``. Delegates to + :meth:`SimulationManager.create_rigid_constraint`. + + Args: + env: The environment instance. + env_ids: Target environment indices. None -> all envs. + obj_a_cfg: SceneEntityCfg pointing at the first RigidObject. + obj_b_cfg: SceneEntityCfg pointing at the second RigidObject. + name: Base constraint name; per-arena names derived by the sim layer. + local_frame_a: Local joint frame on object A. None -> identity (object + A's origin). Accepts (4,4) or (N,4,4). + local_frame_b: Local joint frame on object B. None -> computed per env as + ``inv(pose_B) @ pose_A`` so the constraint welds the objects at their + current relative pose. Accepts (4,4) or (N,4,4). + + Raises: + RuntimeError: If either entity is not a RigidObject. + """ + obj_a = env.sim.get_asset(obj_a_cfg.uid) + obj_b = env.sim.get_asset(obj_b_cfg.uid) + if not isinstance(obj_a, RigidObject) or not isinstance(obj_b, RigidObject): + logger.log_error( + f"Constraint '{name}' requires two RigidObjects, but got " + f"{type(obj_a).__name__} and {type(obj_b).__name__}." + ) + env.sim.create_rigid_constraint( + cfg=RigidConstraintCfg( + name=name, + rigid_object_a_uid=obj_a_cfg.uid, + rigid_object_b_uid=obj_b_cfg.uid, + local_frame_a=local_frame_a, + local_frame_b=local_frame_b, + ), + env_ids=env_ids, + ) + + +def remove_rigid_constraint( + env: EmbodiedEnv, + env_ids: torch.Tensor | None, + name: str, +) -> None: + """Remove the named constraint for the given env_ids. + + Delegates to :meth:`SimulationManager.remove_rigid_constraint`. Idempotent: + warns (via the sim layer) if the constraint is not found. + + Args: + env: The environment instance. + env_ids: Target environment indices. None -> all envs. + name: Base constraint name to remove. + """ + env.sim.remove_rigid_constraint(name, env_ids=env_ids) diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 0b2fe922..4ad97cee 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -997,6 +997,56 @@ def from_dict(cls, init_dict: Dict[str, Any]) -> RigidObjectGroupCfg: return cfg +@configclass +class RigidConstraintCfg: + """Configuration for a fixed constraint between two RigidObjects. + + The constraint binds rigid_object_a's entity[i] to rigid_object_b's entity[i] + within arena[i] (one constraint per arena). + + Args: + name: Base constraint name. Per-arena names are derived as ``f"{name}"`` + (single env) or ``f"{name}_{i}"`` (multi env). + rigid_object_a_uid: UID of the first RigidObject (must exist in the sim). + rigid_object_b_uid: UID of the second RigidObject (must exist in the sim). + local_frame_a: 4x4 joint frame in object A's local coordinates. + ``None`` -> identity (object A's origin). Accepts a single + ``(4, 4)`` matrix (shared by all envs) or an ``(N, 4, 4)`` array + (one frame per env). Defaults to None. + local_frame_b: 4x4 joint frame in object B's local coordinates. + ``None`` -> the frame is computed per env as ``inv(pose_B) @ pose_A`` + from the objects' current poses, so the constraint welds the objects + at their *current* relative pose (rather than pulling their origins + together). An explicit ``(4, 4)`` or ``(N, 4, 4)`` value is used + verbatim. Defaults to None. + constraint_type: Reserved for future typed constraints (prismatic, + revolute, spherical, d6). Only ``"fixed"`` is supported in v1. + + .. attention:: + Both objects must be :class:`RigidObject` instances and must share the + same number of arenas. + """ + + name: str = MISSING + """Base name of the constraint (per-arena names are derived from this).""" + + rigid_object_a_uid: str = MISSING + """UID of the first RigidObject.""" + + rigid_object_b_uid: str = MISSING + """UID of the second RigidObject.""" + + local_frame_a: np.ndarray | None = None + """Local joint frame on object A. None -> identity (object A's origin).""" + + local_frame_b: np.ndarray | None = None + """Local joint frame on object B. None -> ``inv(pose_B) @ pose_A`` per env + (weld at the objects' current relative pose).""" + + constraint_type: Literal["fixed"] = "fixed" + """Constraint type. Only ``"fixed"`` is supported in v1.""" + + @configclass class URDFCfg: """Standalone configuration class for URDF assembly.""" diff --git a/embodichain/lab/sim/objects/__init__.py b/embodichain/lab/sim/objects/__init__.py index 7cde4bbf..2254f100 100644 --- a/embodichain/lab/sim/objects/__init__.py +++ b/embodichain/lab/sim/objects/__init__.py @@ -27,6 +27,7 @@ from .robot import Robot, RobotCfg from .light import Light, LightCfg from .gizmo import Gizmo +from .constraint import RigidConstraint from dexsim.engine import RenderBody diff --git a/embodichain/lab/sim/objects/constraint.py b/embodichain/lab/sim/objects/constraint.py new file mode 100644 index 00000000..0d8951bb --- /dev/null +++ b/embodichain/lab/sim/objects/constraint.py @@ -0,0 +1,149 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Rigid constraint wrapper binding two RigidObjects across arenas.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +if TYPE_CHECKING: + from embodichain.lab.sim.cfg import RigidConstraintCfg + from embodichain.lab.sim.objects.rigid_object import RigidObject + +__all__ = ["RigidConstraint"] + + +@dataclass +class RigidConstraint: + """Batch of fixed constraints linking two :class:`RigidObject` instances. + + Each entry binds ``rigid_object_a._entities[i]`` to + ``rigid_object_b._entities[i]`` within arena ``i`` via a dexsim + ``FixedConstraint``. The ``constraint_handles`` list has length + ``num_envs`` with ``None`` wherever the constraint is not active in that + arena, so arena index always equals list index. + + Args: + cfg: The constraint configuration. + constraint_handles: Per-arena dexsim constraint handles (None where inactive). + rigid_object_a: The first RigidObject (None only until set at creation). + rigid_object_b: The second RigidObject (None only until set at creation). + device: The torch device. + """ + + cfg: RigidConstraintCfg + constraint_handles: list[Any] = field(default_factory=list) + rigid_object_a: RigidObject | None = None + rigid_object_b: RigidObject | None = None + device: torch.device = field(default_factory=torch.device("cpu")) + + @property + def num_envs(self) -> int: + """Number of arenas covered by this constraint.""" + return len(self.constraint_handles) + + def get_name(self, env_id: int) -> str: + """Get the per-arena constraint name. + + For single-env constraints, returns the base name. For multi-env + constraints, returns ``f"{base}_{env_id}"``. + + Args: + env_id: The arena index. + + Returns: + The constraint name registered in that arena. + """ + if self.num_envs <= 1: + return self.cfg.name + return f"{self.cfg.name}_{env_id}" + + def _active_env_ids(self, env_ids: Sequence[int] | None) -> list[int]: + """Resolve the requested env_ids, skipping handles that are None.""" + if env_ids is None: + env_ids = range(self.num_envs) + return [i for i in env_ids if self.constraint_handles[i] is not None] + + def get_relative_transform( + self, env_ids: Sequence[int] | None = None + ) -> list[np.ndarray]: + """Get the relative transform of B in A for each active env. + + Args: + env_ids: Subset of arenas. None -> all arenas. Inactive (None) + handles are skipped. + + Returns: + A list of 4x4 numpy arrays, one per active env. + """ + results = [] + for i in self._active_env_ids(env_ids): + results.append(self.constraint_handles[i].get_relative_transform()) + return results + + def get_local_pose( + self, actor_index: int, env_ids: Sequence[int] | None = None + ) -> list[np.ndarray]: + """Get the local pose of the constraint frame for the given actor. + + Args: + actor_index: 0 for object A, 1 for object B. + env_ids: Subset of arenas. None -> all. Inactive handles skipped. + + Returns: + A list of 4x4 numpy arrays, one per active env. + """ + results = [] + for i in self._active_env_ids(env_ids): + results.append(self.constraint_handles[i].get_local_pose(actor_index)) + return results + + def is_valid(self, env_ids: Sequence[int] | None = None) -> list[bool]: + """Check validity of each active constraint handle. + + Args: + env_ids: Subset of arenas. None -> all. Inactive handles skipped. + + Returns: + A list of bools, one per active env. + """ + return [ + self.constraint_handles[i].is_valid() for i in self._active_env_ids(env_ids) + ] + + def destroy( + self, + env_ids: Sequence[int] | None = None, + arena_resolver: Callable[[int], Any] | None = None, + ) -> None: + """Remove this constraint from the specified arenas. + + Args: + env_ids: Subset of arenas to clear. None -> all active arenas. + arena_resolver: Callable returning the arena for a given env index. + Required to actually remove constraints from dexsim. + """ + for i in self._active_env_ids(env_ids): + if arena_resolver is not None: + arena = arena_resolver(i) + arena.remove_constraint(self.get_name(i)) + self.constraint_handles[i] = None diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index e8ac2916..46340601 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -64,6 +64,7 @@ Articulation, Robot, Light, + RigidConstraint, ) from embodichain.lab.sim.objects.gizmo import Gizmo from embodichain.lab.sim.sensors import ( @@ -86,10 +87,11 @@ RigidObjectGroupCfg, ArticulationCfg, RobotCfg, + RigidConstraintCfg, ) from embodichain.lab.sim import VisualMaterial, VisualMaterialCfg from embodichain.utils import configclass, logger -from embodichain.utils.math import look_at_to_pose +from embodichain.utils.math import look_at_to_pose, pose_inv __all__ = [ "SimulationManager", @@ -283,6 +285,7 @@ def __init__( self._markers: Dict[str, MeshObject] = dict() self._rigid_objects: Dict[str, RigidObject] = dict() + self._constraints: Dict[str, RigidConstraint] = dict() self._rigid_object_groups: Dict[str, RigidObjectGroup] = dict() self._soft_objects: Dict[str, SoftObject] = dict() self._cloth_objects: Dict[str, ClothObject] = dict() @@ -435,6 +438,7 @@ def asset_uids(self) -> List[str]: uid_list.extend(list(self._soft_objects.keys())) uid_list.extend(list(self._cloth_objects.keys())) uid_list.extend(list(self._articulations.keys())) + uid_list.extend(list(self._constraints.keys())) return uid_list def _convert_sim_config( @@ -1004,6 +1008,189 @@ def get_rigid_object_uid_list(self) -> List[str]: """ return list(self._rigid_objects.keys()) + @staticmethod + def _broadcast_frame( + frame: np.ndarray | None, + num_envs: int, + env_ids: Sequence[int], + name: str, + ) -> list[np.ndarray]: + """Broadcast a local-frame spec to one matrix per target env. + + Args: + frame: None -> identity; (4,4) -> repeated; (N,4,4) -> indexed per env. + num_envs: Total number of arenas (used to validate (N,4,4)). + env_ids: Target env indices to produce frames for. + name: Constraint name (for error messages). + + Returns: + A list of (4,4) numpy arrays, one per env in env_ids. + + Raises: + RuntimeError: If an (N,4,4) frame's N != num_envs, or shape is invalid. + """ + if frame is None: + identity = np.eye(4, dtype=np.float32) + return [identity for _ in env_ids] + frame_np = np.asarray(frame, dtype=np.float32) + if frame_np.shape == (4, 4): + return [frame_np for _ in env_ids] + if frame_np.ndim == 3 and frame_np.shape[1:] == (4, 4): + if frame_np.shape[0] != num_envs: + logger.log_error( + f"Constraint '{name}' local frame has shape {frame_np.shape} " + f"but num_envs is {num_envs}. Expected ({num_envs}, 4, 4)." + ) + return [frame_np[i] for i in env_ids] + logger.log_error( + f"Constraint '{name}' local frame has invalid shape {frame_np.shape}. " + "Expected None, (4, 4), or (N, 4, 4)." + ) + + @staticmethod + def _normalize_env_ids( + env_ids: Sequence[int] | torch.Tensor | None, + num_envs: int, + ) -> list[int]: + """Normalize an ``env_ids`` spec to a plain ``list[int]``. + + Accepts ``None`` (-> all envs), a ``torch.Tensor`` (as passed by the + :class:`EventManager`), or any ``Sequence[int]``, and returns a list of + Python ints. Normalizing here keeps the per-arena constraint names clean + (e.g. ``"weld_0"`` rather than relying on a tensor's string form) and + avoids depending on implicit tensor-to-int conversions downstream. + + Args: + env_ids: None, a tensor, or a sequence of ints. + num_envs: Total number of arenas (used when env_ids is None). + + Returns: + A list of int env indices. + """ + if env_ids is None: + return list(range(num_envs)) + if isinstance(env_ids, torch.Tensor): + return env_ids.detach().cpu().tolist() + return [int(i) for i in env_ids] + + def create_rigid_constraint( + self, + cfg: RigidConstraintCfg, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> RigidConstraint: + """Create a fixed constraint between two RigidObjects. + + Binds ``rigid_object_a``'s entity[i] to ``rigid_object_b``'s entity[i] + within arena[i], for each env in ``env_ids``. Local frames default to + welding the objects at their *current* relative pose: + ``local_frame_a`` defaults to identity (object A's origin) and + ``local_frame_b`` defaults to ``inv(pose_B) @ pose_A`` (computed per env), + so the offset is preserved rather than the two origins being pulled + together. Pass explicit frames to define a specific joint frame. + + Args: + cfg: The constraint configuration. + env_ids: Target environment indices. Accepts a tensor (as passed by + the :class:`EventManager`) or a sequence of ints. None -> all arenas. + + Returns: + The created :class:`RigidConstraint`. + + Raises: + RuntimeError: If either object is missing, the name is already in use, + a frame shape is invalid, or dexsim fails to create a handle. + """ + # validate constraint type (only fixed supported in v1) + if cfg.constraint_type != "fixed": + logger.log_error( + f"Constraint '{cfg.name}' has unsupported type " + f"'{cfg.constraint_type}'. Only 'fixed' is supported in v1." + ) + + # resolve objects + if cfg.rigid_object_a_uid not in self._rigid_objects: + logger.log_error( + f"RigidObject '{cfg.rigid_object_a_uid}' not found for constraint " + f"'{cfg.name}'. Available: {list(self._rigid_objects.keys())}." + ) + if cfg.rigid_object_b_uid not in self._rigid_objects: + logger.log_error( + f"RigidObject '{cfg.rigid_object_b_uid}' not found for constraint " + f"'{cfg.name}'. Available: {list(self._rigid_objects.keys())}." + ) + rigid_object_a = self._rigid_objects[cfg.rigid_object_a_uid] + rigid_object_b = self._rigid_objects[cfg.rigid_object_b_uid] + + # validate duplicate name + if cfg.name in self._constraints: + logger.log_error( + f"Constraint '{cfg.name}' already exists. Remove it before recreating." + ) + + # validate object entity counts match num_envs + num_envs = self.num_envs + if rigid_object_a.num_instances != num_envs: + logger.log_error( + f"RigidObject '{cfg.rigid_object_a_uid}' has " + f"{rigid_object_a.num_instances} instances but num_envs is {num_envs}." + ) + if rigid_object_b.num_instances != num_envs: + logger.log_error( + f"RigidObject '{cfg.rigid_object_b_uid}' has " + f"{rigid_object_b.num_instances} instances but num_envs is {num_envs}." + ) + + # resolve target env_ids (accepts None / tensor / sequence) + target_env_ids = self._normalize_env_ids(env_ids, num_envs) + + # broadcast local frames. + # local_frame_a defaults to identity (object A's origin). + # local_frame_b defaults to the current relative pose of A w.r.t. B + # (inv(pose_B) @ pose_A), so that with both frames left as None the + # constraint welds the objects at their *current* relative pose instead + # of pulling their origins together. + frames_a = self._broadcast_frame( + cfg.local_frame_a, num_envs, target_env_ids, cfg.name + ) + if cfg.local_frame_b is None: + pose_a = rigid_object_a.get_local_pose(to_matrix=True) + pose_b = rigid_object_b.get_local_pose(to_matrix=True) + frame_b = torch.bmm(pose_inv(pose_b), pose_a) # (N, 4, 4) + frame_b = frame_b.cpu().numpy().astype(np.float32) + frames_b = [frame_b[i] for i in target_env_ids] + else: + frames_b = self._broadcast_frame( + cfg.local_frame_b, num_envs, target_env_ids, cfg.name + ) + + # pre-size handles list with None, fill target envs + handles: list = [None] * num_envs + for idx, env_id in enumerate(target_env_ids): + arena = self.get_env(env_id) + name_i = cfg.name if num_envs <= 1 else f"{cfg.name}_{env_id}" + handle = arena.create_fixed_constraint( + name_i, + rigid_object_a._entities[env_id], + rigid_object_b._entities[env_id], + frames_a[idx], + frames_b[idx], + ) + if handle is None: + logger.log_error( + f"Failed to create constraint '{name_i}' in arena {env_id}." + ) + handles[env_id] = handle + + constraint = RigidConstraint( + cfg=cfg, + constraint_handles=handles, + rigid_object_a=rigid_object_a, + rigid_object_b=rigid_object_b, + device=self.device, + ) + self._constraints[cfg.name] = constraint + return constraint + def get_soft_object_uid_list(self) -> List[str]: """Get current soft body uid list @@ -1020,6 +1207,61 @@ def get_cloth_object_uid_list(self) -> List[str]: """ return list(self._cloth_objects.keys()) + def remove_rigid_constraint( + self, + name: str, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> bool: + """Remove a rigid constraint by name. + + With ``env_ids=None`` the constraint is removed from every arena and + dropped from the registry. With a subset, only those arenas are cleared; + the registry entry is kept until all handles become None. + + Args: + name: The base constraint name. + env_ids: Subset of arenas to clear. Accepts a tensor (as passed by + the :class:`EventManager`) or a sequence of ints. None -> all. + + Returns: + True if the constraint was found (and removed or partially removed), + False if the name is unknown. + """ + constraint = self._constraints.get(name, None) + if constraint is None: + logger.log_warning(f"Constraint '{name}' not found. Nothing to remove.") + return False + + target_env_ids = self._normalize_env_ids(env_ids, constraint.num_envs) + constraint.destroy(env_ids=target_env_ids, arena_resolver=self.get_env) + + # drop from registry if no handles remain active + if all(h is None for h in constraint.constraint_handles): + del self._constraints[name] + return True + + def get_rigid_constraint(self, name: str) -> RigidConstraint | None: + """Get a rigid constraint by its base name. + + Args: + name: The base constraint name. + + Returns: + The constraint, or None if not found. + """ + if name not in self._constraints: + logger.log_warning(f"Constraint '{name}' not found.") + return None + return self._constraints[name] + + def get_rigid_constraint_uid_list(self) -> List[str]: + """Get the list of registered constraint base names. + + Returns: + List[str]: list of constraint names. + """ + return list(self._constraints.keys()) + def add_rigid_object_group(self, cfg: RigidObjectGroupCfg) -> RigidObjectGroup: """Add a rigid object group to the scene. @@ -2271,6 +2513,7 @@ def _sever_wrapper_refs(obj_registry): _sever_wrapper_refs("_gizmos") _sever_wrapper_refs("_markers") _sever_wrapper_refs("_rigid_objects") + _sever_wrapper_refs("_constraints") _sever_wrapper_refs("_rigid_object_groups") _sever_wrapper_refs("_soft_objects") _sever_wrapper_refs("_cloth_objects") @@ -2293,6 +2536,7 @@ def _sever_wrapper_refs(obj_registry): self._arenas.clear() self._markers.clear() self._gizmos.clear() + self._constraints.clear() SimulationManager.reset(self.instance_id) diff --git a/scripts/tutorials/sim/create_rigid_constraint.py b/scripts/tutorials/sim/create_rigid_constraint.py new file mode 100644 index 00000000..2ac63a22 --- /dev/null +++ b/scripts/tutorials/sim/create_rigid_constraint.py @@ -0,0 +1,184 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +""" +This script demonstrates how to attach two rigid objects via a fixed constraint, +observe the constraint holding their relative pose, and then remove it. +""" + +from __future__ import annotations + +import argparse +import sys + +import numpy as np + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim.cfg import ( + RigidObjectCfg, + RigidConstraintCfg, + RigidBodyAttributesCfg, + RenderCfg, +) +from embodichain.lab.sim.shapes import CubeCfg + +# Number of physics sub-steps per update call. +STEPS_PER_UPDATE = 1 +# Print the relative pose every N update calls. +PRINT_EVERY = 20 +# How long to simulate while attached / detached (in update calls). +PHASE_STEPS = 120 + + +def main(): + """Main function to create and run the constraint tutorial scene.""" + + # Parse command line arguments (adds --headless, --num_envs, --device, ...). + parser = argparse.ArgumentParser( + description="Attach and detach two cubes via a fixed rigid constraint" + ) + add_env_launcher_args_to_parser(parser) + args = parser.parse_args() + + # The simulation teardown (``SimulationManager.destroy``) calls ``os._exit``, + # which skips flushing Python's stdout buffer. Line-buffer stdout so every + # ``print`` below is visible even when the script is piped to a file. + sys.stdout.reconfigure(line_buffering=True) + + # Configure the simulation. + sim_cfg = SimulationManagerCfg( + width=1920, + height=1080, + headless=True, + physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) + sim_device=args.device, + render_cfg=RenderCfg(renderer=args.renderer), + num_envs=args.num_envs, + arena_space=3.0, + ) + + sim = SimulationManager(sim_cfg) + + # Shared physics attributes for the two cubes. + physics_attrs = RigidBodyAttributesCfg( + mass=0.2, + dynamic_friction=0.5, + static_friction=0.5, + restitution=0.1, + ) + + # Add two dynamic cubes to the scene. cube_a starts higher than cube_b so + # that, once detached, the lower cube lands first and the relative pose + # visibly changes (while welded, the constraint holds it constant). + cube_a = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="cube_a", + shape=CubeCfg(size=[0.16, 0.16, 0.16]), + attrs=physics_attrs, + init_pos=[0.0, 0.0, 1.40], + ) + ) + cube_b = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="cube_b", + shape=CubeCfg(size=[0.16, 0.16, 0.16]), + attrs=physics_attrs, + init_pos=[0.0, 0.0, 1.20], + ) + ) + + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + + print("[INFO]: Scene setup complete with two cubes (cube_a, cube_b).") + + # --- Phase 1: attach the two cubes with a fixed constraint --------------- + # With default (None) local frames the constraint welds the cubes at their + # *current* relative pose: local_frame_a defaults to identity and + # local_frame_b is computed as inv(pose_B) @ pose_A, so the offset is + # preserved rather than the two origins being pulled together. + constraint = sim.create_rigid_constraint( + cfg=RigidConstraintCfg( + name="cube_weld", + rigid_object_a_uid="cube_a", + rigid_object_b_uid="cube_b", + ) + ) + print("[INFO]: Created constraint 'cube_weld' between cube_a and cube_b.") + + # Open the viewer (unless --headless) so the welded motion is visible. + if not args.headless: + sim.open_window() + + print("[INFO]: Stepping physics while ATTACHED (relative pose held):") + _run_phase(sim, cube_a, cube_b, attached=True) + + # --- Phase 2: remove the constraint ------------------------------------ + sim.remove_rigid_constraint("cube_weld") + assert sim.get_rigid_constraint("cube_weld") is None + print("\n[INFO]: Removed constraint 'cube_weld'. cube_a and cube_b are now free.") + + import time + + time.sleep(2.0) # Wait a moment so the viewer can show the constraint removal. + + print("[INFO]: Stepping physics while DETACHED (relative pose may drift):") + _run_phase(sim, cube_a, cube_b, attached=False) + + print("\n[INFO]: Tutorial complete.") + sim.destroy() + + +def _relative_z(cube_a, cube_b) -> float: + """Return the z-component of cube_b's pose relative to cube_a (env 0). + + This reads the two bodies' world poses directly, so it works both while the + constraint is active (the value stays constant) and after removal (the + value drifts as the cubes move independently). + + Args: + cube_a: The first :class:`RigidObject`. + cube_b: The second :class:`RigidObject`. + + Returns: + The relative z (cube_b.z - cube_a.z) in meters. + """ + pose_a = cube_a.get_local_pose(to_matrix=True) + pose_b = cube_b.get_local_pose(to_matrix=True) + return float(pose_b[0, 2, 3] - pose_a[0, 2, 3]) + + +def _run_phase(sim, cube_a, cube_b, attached: bool) -> None: + """Step the simulation for one phase and print the bodies' relative z. + + Args: + sim: The :class:`SimulationManager`. + cube_a: The first :class:`RigidObject`. + cube_b: The second :class:`RigidObject`. + attached: True while the constraint is active, False after removal. + """ + rel_z = _relative_z(cube_a, cube_b) + print(f" step {0:4d}: relative z (cube_b - cube_a) = {rel_z:.4f} m") + for step in range(1, PHASE_STEPS + 1): + sim.update(step=STEPS_PER_UPDATE) + if step % PRINT_EVERY == 0: + rel_z = _relative_z(cube_a, cube_b) + print(f" step {step:4d}: relative z (cube_b - cube_a) = {rel_z:.4f} m") + + +if __name__ == "__main__": + main() diff --git a/tests/gym/envs/managers/test_event_rigid_constraint.py b/tests/gym/envs/managers/test_event_rigid_constraint.py new file mode 100644 index 00000000..fb40a03e --- /dev/null +++ b/tests/gym/envs/managers/test_event_rigid_constraint.py @@ -0,0 +1,186 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for the rigid-constraint event functors.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from unittest.mock import MagicMock + +from embodichain.utils import configclass +from embodichain.lab.gym.envs.managers.cfg import EventCfg, SceneEntityCfg +from embodichain.lab.gym.envs.managers.event_manager import EventManager +from embodichain.lab.gym.envs.managers.events import ( + create_rigid_constraint, + remove_rigid_constraint, +) + + +class MockRigidObjectForFunctor: + """Stand-in for a RigidObject passing the isinstance check. + + The functor checks ``isinstance(asset, RigidObject)``. To avoid building a + real RigidObject (needs a dexsim World), we monkeypatch the check. + """ + + +def _make_env(obj_a_is_rigid=True, obj_b_is_rigid=True): + """Build a mock env with a spied sim.create/remove_rigid_constraint. + + obj_a / obj_b are instances of MockRigidObjectForFunctor so that, when the + test patches ``events.RigidObject`` to that class, ``isinstance`` passes. + When ``obj_a_is_rigid`` is False, obj_a becomes a plain object (isinstance + fails) to exercise the rejection path. + """ + env = MagicMock() + env.device = torch.device("cpu") + env.num_envs = 4 + + obj_a = MockRigidObjectForFunctor() if obj_a_is_rigid else object() + obj_b = MockRigidObjectForFunctor() if obj_b_is_rigid else object() + env.sim.get_asset.side_effect = lambda uid: {"cube": obj_a, "block": obj_b}[uid] + + env.sim.create_rigid_constraint = MagicMock(return_value=MagicMock()) + env.sim.remove_rigid_constraint = MagicMock(return_value=True) + return env, obj_a, obj_b + + +def test_create_functor_delegates_to_sim(monkeypatch): + """create functor resolves both objects and forwards to sim.create_rigid_constraint.""" + env, obj_a, obj_b = _make_env() + monkeypatch.setattr( + "embodichain.lab.gym.envs.managers.events.RigidObject", + MockRigidObjectForFunctor, + ) + + env_ids = torch.tensor([0, 2]) + create_rigid_constraint( + env, + env_ids, + obj_a_cfg=SceneEntityCfg(uid="cube"), + obj_b_cfg=SceneEntityCfg(uid="block"), + name="weld", + ) + + env.sim.create_rigid_constraint.assert_called_once() + call_kwargs = env.sim.create_rigid_constraint.call_args + assert call_kwargs.kwargs["env_ids"] is env_ids + cfg = call_kwargs.kwargs["cfg"] + assert cfg.name == "weld" + assert cfg.rigid_object_a_uid == "cube" + assert cfg.rigid_object_b_uid == "block" + assert cfg.local_frame_a is None + assert cfg.local_frame_b is None + + +def test_create_functor_forwards_frames(monkeypatch): + """create functor forwards local frames into the cfg.""" + env, _, _ = _make_env() + monkeypatch.setattr( + "embodichain.lab.gym.envs.managers.events.RigidObject", + MockRigidObjectForFunctor, + ) + frame = np.eye(4, dtype=np.float32) + create_rigid_constraint( + env, + None, + obj_a_cfg=SceneEntityCfg(uid="cube"), + obj_b_cfg=SceneEntityCfg(uid="block"), + name="weld", + local_frame_a=frame, + local_frame_b=frame, + ) + cfg = env.sim.create_rigid_constraint.call_args.kwargs["cfg"] + np.testing.assert_allclose(cfg.local_frame_a, frame) + + +def test_create_functor_rejects_non_rigid_object(monkeypatch): + """A non-RigidObject asset raises (log_error raises RuntimeError).""" + # obj_a is a plain object -> isinstance fails after the patch. + env, obj_a, obj_b = _make_env(obj_a_is_rigid=False, obj_b_is_rigid=True) + monkeypatch.setattr( + "embodichain.lab.gym.envs.managers.events.RigidObject", + MockRigidObjectForFunctor, + ) + with pytest.raises(RuntimeError): + create_rigid_constraint( + env, + None, + obj_a_cfg=SceneEntityCfg(uid="cube"), + obj_b_cfg=SceneEntityCfg(uid="block"), + name="weld", + ) + env.sim.create_rigid_constraint.assert_not_called() + + +def test_remove_functor_delegates(): + """remove functor forwards name + env_ids to sim.remove_rigid_constraint.""" + env, _, _ = _make_env() + env_ids = torch.tensor([1, 3]) + remove_rigid_constraint(env, env_ids, name="weld") + env.sim.remove_rigid_constraint.assert_called_once_with("weld", env_ids=env_ids) + + +def test_remove_functor_none_env_ids(): + """remove functor forwards env_ids=None correctly.""" + env, _, _ = _make_env() + remove_rigid_constraint(env, None, name="weld") + env.sim.remove_rigid_constraint.assert_called_once_with("weld", env_ids=None) + + +@configclass +class _AttachEventsCfg: + attach: EventCfg = EventCfg( + func=create_rigid_constraint, + mode="attach", + params={ + "obj_a_cfg": SceneEntityCfg(uid="cube"), + "obj_b_cfg": SceneEntityCfg(uid="block"), + "name": "weld", + }, + ) + + +def test_custom_mode_apply_invokes_functor_with_env_ids(monkeypatch): + """EventManager.apply(mode="attach", env_ids) calls the functor once with those env_ids.""" + # Build a minimal env stand-in that EventManager needs: num_envs, device, sim. + env = MagicMock() + env.num_envs = 4 + env.device = torch.device("cpu") + env.sim = MagicMock() + # SceneEntityCfg.resolve() checks scene.asset_uids for each uid; list them. + env.sim.asset_uids = ["cube", "block"] + # get_asset returns MockRigidObjectForFunctor instances so isinstance passes. + obj_a = MockRigidObjectForFunctor() + obj_b = MockRigidObjectForFunctor() + env.sim.get_asset.side_effect = lambda uid: {"cube": obj_a, "block": obj_b}[uid] + env.sim.create_rigid_constraint = MagicMock() + + monkeypatch.setattr( + "embodichain.lab.gym.envs.managers.events.RigidObject", + MockRigidObjectForFunctor, + ) + + manager = EventManager(cfg=_AttachEventsCfg(), env=env) + + env_ids = torch.tensor([0, 1]) + manager.apply(mode="attach", env_ids=env_ids) + + env.sim.create_rigid_constraint.assert_called_once() + assert env.sim.create_rigid_constraint.call_args.kwargs["env_ids"] is env_ids diff --git a/tests/sim/objects/test_rigid_constraint.py b/tests/sim/objects/test_rigid_constraint.py new file mode 100644 index 00000000..55bd9d43 --- /dev/null +++ b/tests/sim/objects/test_rigid_constraint.py @@ -0,0 +1,589 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for the RigidConstraint sim-layer wrapper and its config.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from unittest.mock import MagicMock + +from dataclasses import MISSING + +from embodichain.lab.sim.cfg import RigidConstraintCfg +from embodichain.lab.sim.objects.constraint import RigidConstraint +from embodichain.lab.sim.sim_manager import SimulationManager + + +def test_rigid_constraint_cfg_defaults(): + """RigidConstraintCfg requires name + both object uids; frames default None.""" + cfg = RigidConstraintCfg( + name="weld", + rigid_object_a_uid="cube", + rigid_object_b_uid="block", + ) + assert cfg.name == "weld" + assert cfg.rigid_object_a_uid == "cube" + assert cfg.rigid_object_b_uid == "block" + assert cfg.local_frame_a is None + assert cfg.local_frame_b is None + assert cfg.constraint_type == "fixed" + + +def test_rigid_constraint_cfg_required_fields_are_missing(): + """Required fields default to the MISSING sentinel.""" + assert RigidConstraintCfg.__dataclass_fields__["name"].default is MISSING + assert ( + RigidConstraintCfg.__dataclass_fields__["rigid_object_a_uid"].default is MISSING + ) + assert ( + RigidConstraintCfg.__dataclass_fields__["rigid_object_b_uid"].default is MISSING + ) + + +def test_rigid_constraint_cfg_accepts_frames(): + """Local frames accept 4x4 numpy arrays.""" + frame = np.eye(4, dtype=np.float32) + cfg = RigidConstraintCfg( + name="weld", + rigid_object_a_uid="cube", + rigid_object_b_uid="block", + local_frame_a=frame, + local_frame_b=frame, + ) + np.testing.assert_allclose(cfg.local_frame_a, frame) + + +def _make_handle(name="weld_0", rel_z=0.0, valid=True): + """Build a mock dexsim constraint handle.""" + h = MagicMock() + h.get_name.return_value = name + h.is_valid.return_value = valid + rel = np.eye(4, dtype=np.float32) + rel[2, 3] = rel_z + h.get_relative_transform.return_value = rel + h.get_local_pose.return_value = np.eye(4, dtype=np.float32) + return h + + +def _make_rigid_object(uid="cube", num_envs=4): + """Build a mock RigidObject with a per-arena entity list.""" + obj = MagicMock() + obj.uid = uid + obj.num_instances = num_envs + obj._entities = [MagicMock() for _ in range(num_envs)] + return obj + + +def test_rigid_constraint_num_envs_and_init(): + """RigidConstraint exposes num_envs and stores handles + object refs.""" + cfg = RigidConstraintCfg( + name="weld", + rigid_object_a_uid="cube", + rigid_object_b_uid="block", + ) + handles = [_make_handle("weld_0"), None, _make_handle("weld_2"), None] + obj_a = _make_rigid_object("cube", 4) + obj_b = _make_rigid_object("block", 4) + + constraint = RigidConstraint( + cfg=cfg, + constraint_handles=handles, + rigid_object_a=obj_a, + rigid_object_b=obj_b, + device=torch.device("cpu"), + ) + assert constraint.num_envs == 4 + assert constraint.rigid_object_a is obj_a + assert constraint.rigid_object_b is obj_b + assert len(constraint.constraint_handles) == 4 + + +def test_rigid_constraint_get_name_single_and_multi_env(): + """Single env keeps the base name; multi env appends the arena index.""" + cfg_single = RigidConstraintCfg( + name="weld", + rigid_object_a_uid="cube", + rigid_object_b_uid="block", + ) + c_single = RigidConstraint( + cfg_single, + [_make_handle("weld")], + MagicMock(), + MagicMock(), + torch.device("cpu"), + ) + assert c_single.get_name(0) == "weld" + + cfg_multi = RigidConstraintCfg( + name="weld", + rigid_object_a_uid="cube", + rigid_object_b_uid="block", + ) + handles = [_make_handle("weld_0"), _make_handle("weld_1")] + c_multi = RigidConstraint( + cfg_multi, handles, MagicMock(), MagicMock(), torch.device("cpu") + ) + assert c_multi.get_name(0) == "weld_0" + assert c_multi.get_name(1) == "weld_1" + + +def test_rigid_constraint_get_relative_transform_skips_none(): + """get_relative_transform skips None handles and only returns for active envs.""" + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="a", rigid_object_b_uid="b" + ) + handles = [ + _make_handle("weld_0", rel_z=0.1), + None, + _make_handle("weld_2", rel_z=0.2), + None, + ] + constraint = RigidConstraint( + cfg, handles, MagicMock(), MagicMock(), torch.device("cpu") + ) + + # default: all env_ids, skips None + transforms = constraint.get_relative_transform() + assert len(transforms) == 2 + assert transforms[0][2, 3] == pytest.approx(0.1) + assert transforms[1][2, 3] == pytest.approx(0.2) + + # explicit subset including a None handle is skipped + transforms_subset = constraint.get_relative_transform(env_ids=[1, 2]) + assert len(transforms_subset) == 1 + assert transforms_subset[0][2, 3] == pytest.approx(0.2) + + +def test_rigid_constraint_is_valid(): + """is_valid reports per-env validity, skipping None handles.""" + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="a", rigid_object_b_uid="b" + ) + handles = [_make_handle(valid=True), None, _make_handle(valid=False), None] + constraint = RigidConstraint( + cfg, handles, MagicMock(), MagicMock(), torch.device("cpu") + ) + assert constraint.is_valid() == [True, False] + + +def test_rigid_constraint_destroy_calls_arena_remove_per_env(): + """destroy calls arena.remove_constraint for each active handle in env_ids.""" + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="a", rigid_object_b_uid="b" + ) + handles = [ + _make_handle("weld_0"), + None, + _make_handle("weld_2"), + _make_handle("weld_3"), + ] + constraint = RigidConstraint( + cfg, handles, MagicMock(), MagicMock(), torch.device("cpu") + ) + + arenas = [MagicMock() for _ in range(4)] + arena_resolver = lambda i: arenas[i] + + constraint.destroy(env_ids=[0, 2], arena_resolver=arena_resolver) + arenas[0].remove_constraint.assert_called_once_with("weld_0") + arenas[2].remove_constraint.assert_called_once_with("weld_2") + arenas[1].remove_constraint.assert_not_called() + arenas[3].remove_constraint.assert_not_called() + # cleared handles become None + assert constraint.constraint_handles[0] is None + assert constraint.constraint_handles[2] is None + assert constraint.constraint_handles[3] is not None # not in env_ids + + +def test_rigid_constraint_destroy_all_returns_all_cleared(): + """destroy with env_ids=None clears every active handle.""" + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="a", rigid_object_b_uid="b" + ) + handles = [_make_handle("weld_0"), None, _make_handle("weld_2"), None] + constraint = RigidConstraint( + cfg, handles, MagicMock(), MagicMock(), torch.device("cpu") + ) + arenas = [MagicMock() for _ in range(4)] + constraint.destroy(env_ids=None, arena_resolver=lambda i: arenas[i]) + assert all(h is None for h in constraint.constraint_handles) + + +class MockArena: + """Mock dexsim arena that records created constraints.""" + + def __init__(self, fail_indices=None): + self.created = [] # list of (name, actor0, actor1, frame_a, frame_b) + self.removed = [] # list of names + self.fail_indices = set(fail_indices or []) + + def create_fixed_constraint(self, name, actor0, actor1, local_frame0, local_frame1): + self.created.append((name, actor0, actor1, local_frame0, local_frame1)) + if len(self.created) - 1 in self.fail_indices: + return None + h = MagicMock() + h.get_name.return_value = name + h.is_valid.return_value = True + h.get_relative_transform.return_value = np.eye(4, dtype=np.float32) + return h + + def remove_constraint(self, name): + self.removed.append(name) + + +class _RigidConstraintTestSim: + """A SimulationManager stand-in exposing only the constraint registry path. + + We avoid constructing a real dexsim World (which needs a GPU/window). Instead + we drive create_rigid_constraint by giving it a fake `self` with the + attributes the method touches: _rigid_objects, _arenas/_env, num_envs, device. + """ + + def __init__(self, num_envs=4, arenas=None): + self._rigid_objects = {} + self._constraints = {} + self.device = torch.device("cpu") + if num_envs == 1: + self._arenas = [] + self._env = arenas[0] if arenas else MockArena() + else: + self._arenas = arenas or [MockArena() for _ in range(num_envs)] + self._env = None + + @property + def num_envs(self): + return len(self._arenas) if self._arenas else 1 + + def get_env(self, arena_index=-1): + if arena_index >= 0 and self._arenas: + return self._arenas[arena_index] + return self._env + + # bind the real method under test + create_rigid_constraint = SimulationManager.create_rigid_constraint + remove_rigid_constraint = SimulationManager.remove_rigid_constraint + get_rigid_constraint = SimulationManager.get_rigid_constraint + get_rigid_constraint_uid_list = SimulationManager.get_rigid_constraint_uid_list + _broadcast_frame = staticmethod(SimulationManager._broadcast_frame) + _normalize_env_ids = staticmethod(SimulationManager._normalize_env_ids) + + +def _register_object(sim, uid, num_envs, z=0.0): + """Register a mock RigidObject. + + ``get_local_pose`` returns a real ``(num_envs, 4, 4)`` tensor so the + constraint's default ``local_frame_b`` computation (which reads both + objects' current poses) works under the mock. ``z`` sets the per-env + translation so two objects can be placed at different heights. + """ + obj = MagicMock() + obj.uid = uid + obj.num_instances = num_envs + obj._entities = [MagicMock(name=f"{uid}_{i}") for i in range(num_envs)] + pose = torch.eye(4, dtype=torch.float32).unsqueeze(0).repeat(num_envs, 1, 1) + pose[:, 2, 3] = z + obj.get_local_pose.return_value = pose + sim._rigid_objects[uid] = obj + return obj + + +def test_create_rigid_constraint_resolves_both_objects_all_envs(): + """create builds one handle per arena and stores a RigidConstraint.""" + sim = _RigidConstraintTestSim(num_envs=4) + _register_object(sim, "cube", 4) + _register_object(sim, "block", 4) + + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + constraint = sim.create_rigid_constraint(cfg) + + assert cfg.name in sim._constraints + assert constraint.num_envs == 4 + assert all(h is not None for h in constraint.constraint_handles) + # each arena got exactly one create call with the right actors + for i, arena in enumerate(sim._arenas): + assert arena.created[0][0] == f"weld_{i}" + assert arena.created[0][1] is sim._rigid_objects["cube"]._entities[i] + assert arena.created[0][2] is sim._rigid_objects["block"]._entities[i] + + +def test_create_rigid_constraint_single_env_uses_global_env(): + """Single-env create routes through the global env and keeps the base name.""" + arena = MockArena() + sim = _RigidConstraintTestSim(num_envs=1, arenas=[arena]) + _register_object(sim, "cube", 1) + _register_object(sim, "block", 1) + + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + constraint = sim.create_rigid_constraint(cfg) + assert constraint.constraint_handles[0] is not None + assert arena.created[0][0] == "weld" # base name, no suffix + + +def test_create_rigid_constraint_subset_env_ids(): + """env_ids subset populates only those arenas; others stay None.""" + sim = _RigidConstraintTestSim(num_envs=4) + _register_object(sim, "cube", 4) + _register_object(sim, "block", 4) + + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + constraint = sim.create_rigid_constraint(cfg, env_ids=[0, 2]) + assert constraint.constraint_handles[0] is not None + assert constraint.constraint_handles[1] is None + assert constraint.constraint_handles[2] is not None + assert constraint.constraint_handles[3] is None + # only arenas 0 and 2 got a create call + assert len(sim._arenas[0].created) == 1 + assert len(sim._arenas[1].created) == 0 + assert len(sim._arenas[2].created) == 1 + assert len(sim._arenas[3].created) == 0 + + +def test_create_rigid_constraint_missing_object_raises(): + """A missing object uid raises (log_error raises RuntimeError by default).""" + sim = _RigidConstraintTestSim(num_envs=4) + _register_object(sim, "cube", 4) + # block not registered + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + with pytest.raises(RuntimeError): + sim.create_rigid_constraint(cfg) + + +def test_create_rigid_constraint_duplicate_name_raises(): + """A duplicate base name raises.""" + sim = _RigidConstraintTestSim(num_envs=4) + _register_object(sim, "cube", 4) + _register_object(sim, "block", 4) + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + sim.create_rigid_constraint(cfg) + with pytest.raises(RuntimeError): + sim.create_rigid_constraint(cfg) + + +def test_create_rigid_constraint_failed_handle_raises(): + """If dexsim returns None for a handle, log_error raises.""" + sim = _RigidConstraintTestSim( + num_envs=2, arenas=[MockArena(fail_indices=[0]), MockArena()] + ) + _register_object(sim, "cube", 2) + _register_object(sim, "block", 2) + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + with pytest.raises(RuntimeError): + sim.create_rigid_constraint(cfg) + + +def test_broadcast_frame_none_to_identity(): + """None frame broadcasts to identity per env.""" + sim = _RigidConstraintTestSim(num_envs=3) + frames = sim._broadcast_frame(None, num_envs=3, env_ids=[0, 1, 2], name="weld") + assert len(frames) == 3 + for f in frames: + np.testing.assert_allclose(f, np.eye(4)) + + +def test_broadcast_frame_4x4_repeats(): + """A single 4x4 matrix repeats across all envs.""" + sim = _RigidConstraintTestSim(num_envs=3) + frame = np.eye(4, dtype=np.float32) * 2 + frames = sim._broadcast_frame(frame, num_envs=3, env_ids=[0, 1, 2], name="weld") + assert len(frames) == 3 + for f in frames: + np.testing.assert_allclose(f, frame) + + +def test_broadcast_frame_N4x4_indexes(): + """An (N,4,4) array indexes per env and requires N == num_envs.""" + sim = _RigidConstraintTestSim(num_envs=3) + frames_in = np.stack([np.eye(4) * i for i in range(3)], axis=0).astype(np.float32) + frames = sim._broadcast_frame(frames_in, num_envs=3, env_ids=[0, 1, 2], name="weld") + for i, f in enumerate(frames): + np.testing.assert_allclose(f, frames_in[i]) + + # wrong N raises + bad = np.stack([np.eye(4)] * 2, axis=0).astype(np.float32) + with pytest.raises(RuntimeError): + sim._broadcast_frame(bad, num_envs=3, env_ids=[0, 1, 2], name="weld") + + +def test_normalize_env_ids_handles_tensor_and_none(): + """_normalize_env_ids accepts None / tensor / sequence and returns list[int].""" + assert SimulationManager._normalize_env_ids(None, 3) == [0, 1, 2] + assert SimulationManager._normalize_env_ids(torch.tensor([0, 2]), 4) == [0, 2] + assert SimulationManager._normalize_env_ids([1, 3], 4) == [1, 3] + out = SimulationManager._normalize_env_ids(torch.tensor([0, 1, 2]), 3) + assert isinstance(out, list) + assert all(isinstance(i, int) for i in out) + + +def test_create_rigid_constraint_accepts_tensor_env_ids(): + """A tensor env_ids (as passed by EventManager) yields clean per-arena names. + + Regression for the type-contract mismatch between the functor (which + forwards a torch.Tensor) and the sim API. The per-arena dexsim name must be + ``"weld_0"`` (not a tensor stringification) so create and remove agree. + """ + sim = _RigidConstraintTestSim(num_envs=2) + _register_object(sim, "cube", 2, z=1.4) + _register_object(sim, "block", 2, z=1.2) + + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + sim.create_rigid_constraint(cfg, env_ids=torch.tensor([0, 1])) + for i, arena in enumerate(sim._arenas): + assert arena.created[0][0] == f"weld_{i}" # clean name, no tensor string + + # removal via tensor env_ids must reach the same dexsim name + sim.remove_rigid_constraint("weld", env_ids=torch.tensor([0])) + assert "weld_0" in sim._arenas[0].removed + assert sim._arenas[1].removed == [] # env 1 still attached + + +def test_remove_rigid_constraint_all_envs(): + """remove with env_ids=None clears every arena and drops the registry entry.""" + sim = _RigidConstraintTestSim(num_envs=4) + _register_object(sim, "cube", 4) + _register_object(sim, "block", 4) + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + sim.create_rigid_constraint(cfg) + + removed = sim.remove_rigid_constraint("weld") + assert removed is True + assert "weld" not in sim._constraints + # each arena got remove_constraint with its per-env name + for i, arena in enumerate(sim._arenas): + assert f"weld_{i}" in arena.removed + + +def test_remove_rigid_constraint_subset_keeps_others(): + """remove with a subset env_ids clears only those arenas; registry kept.""" + sim = _RigidConstraintTestSim(num_envs=4) + _register_object(sim, "cube", 4) + _register_object(sim, "block", 4) + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + sim.create_rigid_constraint(cfg) + + removed = sim.remove_rigid_constraint("weld", env_ids=[0, 2]) + assert removed is True + # still in registry because envs 1,3 remain active + assert "weld" in sim._constraints + assert sim._constraints["weld"].constraint_handles[0] is None + assert sim._constraints["weld"].constraint_handles[1] is not None + assert sim._constraints["weld"].constraint_handles[2] is None + assert sim._constraints["weld"].constraint_handles[3] is not None + assert "weld_0" in sim._arenas[0].removed + assert "weld_2" in sim._arenas[2].removed + assert sim._arenas[1].removed == [] + + +def test_remove_rigid_constraint_unknown_name_warns_false(): + """remove on an unknown name returns False without raising.""" + sim = _RigidConstraintTestSim(num_envs=4) + removed = sim.remove_rigid_constraint("nope") + assert removed is False + + +def test_get_rigid_constraint_and_uid_list(): + """get returns the constraint; uid list lists all registered names.""" + sim = _RigidConstraintTestSim(num_envs=2) + _register_object(sim, "cube", 2) + _register_object(sim, "block", 2) + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + sim.create_rigid_constraint(cfg) + assert sim.get_rigid_constraint("weld") is not None + assert sim.get_rigid_constraint("nope") is None + assert sim.get_rigid_constraint_uid_list() == ["weld"] + + +def test_partial_remove_then_all_drops_registry(): + """Subset remove then removing remaining envs drops the registry entry.""" + sim = _RigidConstraintTestSim(num_envs=2) + _register_object(sim, "cube", 2) + _register_object(sim, "block", 2) + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube", rigid_object_b_uid="block" + ) + sim.create_rigid_constraint(cfg) + sim.remove_rigid_constraint("weld", env_ids=[0]) + assert "weld" in sim._constraints + sim.remove_rigid_constraint("weld", env_ids=[1]) + assert "weld" not in sim._constraints + + +def test_create_rigid_constraint_default_frame_b_preserves_relative_pose(): + """With local_frame_b=None, frame_b = inv(pose_B) @ pose_A (preserves offset). + + cube_a at z=1.4, cube_b at z=1.2 -> B is 0.2 below A. The computed + local_frame_b must translate +0.2 in z (A's pose relative to B) so the + constraint welds the cubes at their current relative pose instead of + pulling their origins together. + """ + sim = _RigidConstraintTestSim(num_envs=2) + _register_object(sim, "cube_a", 2, z=1.4) + _register_object(sim, "cube_b", 2, z=1.2) + + cfg = RigidConstraintCfg( + name="weld", rigid_object_a_uid="cube_a", rigid_object_b_uid="cube_b" + ) + sim.create_rigid_constraint(cfg) + + for i, arena in enumerate(sim._arenas): + # arena.created[0] = (name, actor0, actor1, frame_a, frame_b) + frame_a = arena.created[0][3] + frame_b = arena.created[0][4] + # frame_a defaults to identity + np.testing.assert_allclose(frame_a, np.eye(4), atol=1e-6) + # frame_b = inv(pose_B) @ pose_A = translate(0, 0, +0.2) + np.testing.assert_allclose(frame_b[:3, 3], [0.0, 0.0, 0.2], atol=1e-5) + + +def test_create_rigid_constraint_explicit_frame_b_used_verbatim(): + """An explicit local_frame_b is broadcast verbatim (no pose computation).""" + sim = _RigidConstraintTestSim(num_envs=2) + _register_object(sim, "cube_a", 2, z=1.4) + _register_object(sim, "cube_b", 2, z=1.2) + + explicit = np.eye(4, dtype=np.float32) * 3.0 + cfg = RigidConstraintCfg( + name="weld", + rigid_object_a_uid="cube_a", + rigid_object_b_uid="cube_b", + local_frame_b=explicit, + ) + sim.create_rigid_constraint(cfg) + for arena in sim._arenas: + frame_b = arena.created[0][4] + np.testing.assert_allclose(frame_b, explicit) diff --git a/tests/sim/test_rigid_constraint_integration.py b/tests/sim/test_rigid_constraint_integration.py new file mode 100644 index 00000000..4cf49c51 --- /dev/null +++ b/tests/sim/test_rigid_constraint_integration.py @@ -0,0 +1,168 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Real-sim integration smoke test for rigid constraints. + +Mirrors the dexsim ``test_constraint.py`` contract at the EmbodiChain +:class:`SimulationManager` layer: two dynamic objects welded by a fixed +constraint keep their relative transform under physics, and detaching lets +them separate. + +Skipped when the required asset is not present on disk (it is downloaded on +demand by ``embodichain.data.get_data_path``) or when CUDA is unavailable. +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import ( + RigidObjectCfg, + RigidConstraintCfg, + RigidBodyAttributesCfg, +) +from embodichain.lab.sim.shapes import MeshCfg + +DUCK_PATH = "ToyDuck/toy_duck.glb" + + +def _can_run_sim(device: str) -> bool: + """True only if the device is usable and the test asset is present.""" + if device == "cuda" and not torch.cuda.is_available(): + return False + try: + return os.path.isfile(get_data_path(DUCK_PATH)) + except Exception: + return False + + +class BaseRigidConstraintTest: + """Shared setup for the CPU and CUDA integration tests.""" + + def _delta_z(self) -> float: + """Return duck_b.z - duck_a.z (env 0) from the bodies' current poses.""" + pose_a = self.duck_a.get_local_pose(to_matrix=True) + pose_b = self.duck_b.get_local_pose(to_matrix=True) + return float(pose_b[0, 2, 3] - pose_a[0, 2, 3]) + + def setup_simulation(self, sim_device: str) -> None: + if not _can_run_sim(sim_device): + pytest.skip( + f"Cannot run rigid-constraint integration test on {sim_device}." + ) + config = SimulationManagerCfg(headless=True, sim_device=sim_device, num_envs=1) + self.sim = SimulationManager(config) + self.sim.enable_physics(False) + + duck_path = get_data_path(DUCK_PATH) + # Two dynamic ducks at different heights; with default (None) local + # frames the constraint welds them at their current relative pose. + attrs_a = RigidBodyAttributesCfg() + attrs_a.mass = 0.2 + attrs_b = RigidBodyAttributesCfg() + attrs_b.mass = 0.1 + self.duck_a = self.sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="duck_a", + shape=MeshCfg(fpath=duck_path), + body_type="dynamic", + init_pos=[0.0, 0.0, 1.4], + attrs=attrs_a, + ), + ) + self.duck_b = self.sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="duck_b", + shape=MeshCfg(fpath=duck_path), + body_type="dynamic", + init_pos=[0.0, 0.0, 1.2], + attrs=attrs_b, + ), + ) + + if sim_device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): + self.sim.init_gpu_physics() + self.sim.enable_physics(True) + + def teardown_method(self): + """Destroy the simulation and flush the deferred-cleanup queue. + + Mirrors the teardown pattern in ``tests/sim/objects/test_rigid_object.py``. + ``conftest.py`` sets ``EMBODICHAIN_SIM_EXIT_PROCESS=0`` so ``destroy()`` + does not call ``os._exit`` during tests. Guarded for the skip case where + ``setup_simulation`` never created a ``sim``. + """ + import gc + + if getattr(self, "sim", None) is not None: + self.sim.destroy() + SimulationManager.flush_cleanup_queue() + self.__dict__.clear() + gc.collect() + + def test_fixed_constraint_holds_relative_pose(self): + """Welded objects keep their relative transform; detaching lets them move.""" + constraint = self.sim.create_rigid_constraint( + RigidConstraintCfg( + name="weld", + rigid_object_a_uid="duck_a", + rigid_object_b_uid="duck_b", + ) + ) + assert constraint.is_valid() == [True] + + pose_a = self.duck_a.get_local_pose(to_matrix=True) + pose_b = self.duck_b.get_local_pose(to_matrix=True) + initial_delta_z = float(pose_b[0, 2, 3] - pose_a[0, 2, 3]) + + # Step physics; the constraint should hold the relative transform. + for _ in range(120): + self.sim.update(step=1) + + rel = constraint.get_relative_transform()[0] + np.testing.assert_allclose(rel[:3, 3], np.zeros(3), atol=2e-2) + pose_a2 = self.duck_a.get_local_pose(to_matrix=True) + pose_b2 = self.duck_b.get_local_pose(to_matrix=True) + delta_z = float(pose_b2[0, 2, 3] - pose_a2[0, 2, 3]) + assert abs(delta_z - initial_delta_z) < 2e-2 + + # Detach and confirm the relative pose is no longer held: once free, + # the lower duck (duck_b) lands first and the gap closes, so the + # relative z drifts away from the value the constraint was holding. + self.sim.remove_rigid_constraint("weld") + assert self.sim.get_rigid_constraint("weld") is None + held_delta_z = self._delta_z() + for _ in range(120): + self.sim.update(step=1) + final_delta_z = self._delta_z() + assert abs(final_delta_z - held_delta_z) > 2e-2 + + +class TestRigidConstraintCPU(BaseRigidConstraintTest): + def setup_method(self) -> None: + self.setup_simulation("cpu") + + +@pytest.mark.skip(reason="Skipping CUDA tests temporarily") +class TestRigidConstraintCUDA(BaseRigidConstraintTest): + def setup_method(self) -> None: + self.setup_simulation("cuda")