From 81d1f964b259736432851395f2851455000df364 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:01:51 +0000 Subject: [PATCH 01/16] Add rigid constraint design spec Design for attaching two RigidObjects via a fixed physics constraint and removing it, with a standalone sim-layer API (SimulationManager + RigidConstraint) and an on-demand event functor in events.py. Co-Authored-By: Claude --- .../2026-06-29-rigid-constraint-design.md | 332 ++++++++++++++++++ 1 file changed, 332 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-29-rigid-constraint-design.md diff --git a/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md b/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md new file mode 100644 index 00000000..b1d5c07c --- /dev/null +++ b/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md @@ -0,0 +1,332 @@ +# Rigid Object Constraint — Design Spec + +- **Date:** 2026-06-29 +- **Status:** Approved +- **Topic:** Fixed constraint support between two `RigidObject`s, plus an on-demand event functor + +## 1. Goal + +Add the ability to attach two rigid objects via a fixed physics constraint and to remove that constraint again, both as a standalone simulation API (usable outside the gym) and as an on-demand event functor triggered from a task environment. + +### Functional requirements + +1. Two `RigidObject`s can be attached via a constraint. +2. The constraint between the two rigid objects can be removed. +3. A functor creates and removes the constraint inside a task environment. + +### Non-goals (v1, deferred with extension points) + +- Prismatic / revolute / spherical / d6 typed constraints (reserved `constraint_type` field). +- Per-constraint physics tuning (limits, drive, motion) — comes with typed constraints. +- Articulation ↔ rigid and articulation ↔ articulation constraints — `RigidObject` ↔ `RigidObject` only. +- Auto-detach-on-reset baked into the sim layer — kept as task policy. + +## 2. Reference + +The design mirrors the dexsim `fixed_constraint` example (`dexsim/examples/python/physics/rigidbody/fixed_constraint.py`) and lifts its API onto EmbodiChain's batched object/manager layer. + +dexsim constraint API (bound on `Arena`, inherited by `Env`; C++ in `dexsim/cpp/pybind/environment/environment.cpp`): + +- `arena.create_fixed_constraint(name, actor0, actor1, local_frame0=I, local_frame1=I) -> FixedConstraint | None` +- `arena.remove_constraint(name)` +- `arena.get_constraint(name) -> Constraint | None` +- `arena.get_all_constraints() -> list[Constraint]` +- Constraint handle exposes `get_name()`, `get_constraint_type()`, `is_valid()`, `get_local_pose(idx)`, `get_relative_transform()`. +- Typed variants exist (`create_prismatic_constraint`, `create_revolute_constraint`, `create_spherical_constraint`, `create_d6_constraint`) — reserved for later. + +## 3. Architecture + +Two layers: + +``` +SIM LAYER (standalone, usable without gym) + SimulationManager RigidConstraint (new) + ├─ create_rigid_constraint(...) ──► batch wrapper over N arenas + ├─ remove_rigid_constraint(name) ├─ per-arena dexsim handles + ├─ get_rigid_constraint(name) ├─ obj_a / obj_b refs + └─ self._constraints: dict ├─ get_relative_transform() + {name: RigidConstraint} ├─ get_local_pose(idx) + └─ destroy() + delegates to dexsim Arena.create_fixed_constraint / + Arena.remove_constraint (per env_id / arena) + +FUNCTOR LAYER (gym, on-demand) + events.py + ├─ create_rigid_constraint(env, env_ids, obj_a_cfg, obj_b_cfg, name, ...) + └─ remove_rigid_constraint(env, env_ids, name) + registered under EventCfg(mode=, ...); + task triggers via env.event_manager.apply(mode="attach", env_ids) +``` + +### Principles + +1. **One source of truth** — the sim layer owns the constraint registry and all dexsim calls. The functor is a thin adapter: resolve `SceneEntityCfg` → `RigidObject`, then call the sim API. +2. **Per-arena batch symmetry** — `RigidConstraint` mirrors `RigidObject`: N arenas → N dexsim constraint handles, so `env_id i` ↔ arena `i` ↔ handle `i`. Attach/detach can target a subset of `env_ids`. +3. **Fixed-first, extensible** — v1 wires only `create_fixed_constraint`; `RigidConstraintCfg.constraint_type` is reserved (`"fixed"` default). +4. **Local frames default to identity** → attaches at the objects' *current* relative pose. Caller can pass 4×4 or `(N,4,4)`. + +### Why `SimulationManager`, not `RigidObject`, owns the API + +The sim owns `self._arenas`, `self._env`, and every scene-mutation method (`add_rigid_object`, `remove_asset`, `draw_marker`). A constraint lives *between* two objects, so it belongs to the scene owner, not to either object. This keeps `RigidObject` focused on a single body and the constraint registry co-located with the rigid-object registry. + +## 4. Sim-layer API + +### New file: `embodichain/lab/sim/objects/constraint.py` + +```python +@dataclass +class RigidConstraint: + """Batch of fixed constraints linking two RigidObjects across all arenas. + + Each entry binds rigid_object_a's entity[i] to rigid_object_b's entity[i] + within arena[i] via a dexsim FixedConstraint. + """ + cfg: RigidConstraintCfg + constraint_handles: list[Any] # length == num_envs; None where inactive + rigid_object_a: RigidObject + rigid_object_b: RigidObject + device: torch.device + + @property + def num_envs(self) -> int: ... + + def get_relative_transform(self, env_ids=None) -> list[np.ndarray]: ... + def get_local_pose(self, actor_index: int, env_ids=None) -> list[np.ndarray]: ... + def get_name(self, env_id: int) -> str: ... + def is_valid(self, env_ids=None) -> list[bool]: ... + def destroy(self, env_ids: Sequence[int] | None = None) -> None: ... +``` + +`constraint_handles` is a list of length `num_envs` with `None` wherever the constraint is not active in that arena, so **arena index == list index** always holds. + +### `RigidConstraintCfg` (in `embodichain/lab/sim/cfg.py`) + +```python +@configclass +class RigidConstraintCfg: + name: str = MISSING + rigid_object_a_uid: str = MISSING + rigid_object_b_uid: str = MISSING + local_frame_a: np.ndarray | None = None # None → identity; 4x4 or (N,4,4) + local_frame_b: np.ndarray | None = None + constraint_type: Literal["fixed"] = "fixed" # reserved for typed constraints +``` + +### `SimulationManager` additions + +```python +def create_rigid_constraint( + self, cfg: RigidConstraintCfg, env_ids: Sequence[int] | None = None +) -> RigidConstraint: + """Create a fixed constraint between two RigidObjects (env_ids-aware).""" + +def remove_rigid_constraint( + self, name: str, env_ids: Sequence[int] | None = None +) -> bool: + """Remove by base name; idempotent. env_ids subset clears only those arenas.""" + +def get_rigid_constraint(self, name: str) -> RigidConstraint | None: ... +def get_rigid_constraint_uid_list(self) -> list[str]: ... +``` + +- New registry `self._constraints: Dict[str, RigidConstraint]`. +- `_deferred_destroy` severs `_constraints` like the other registries. + +## 5. Functor layer + +Two function-style event functors in `embodichain/lab/gym/envs/managers/events.py`, following the `add-functor` conventions: signature `(env, env_ids, ...) -> None`, `SceneEntityCfg` for entity refs, `from __future__ import annotations`, `TYPE_CHECKING` guard for `EmbodiedEnv`. + +### `create_rigid_constraint` (attach) + +```python +def create_rigid_constraint( + env, env_ids, + 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.""" + obj_a = env.sim.get_asset(obj_a_cfg.uid) + obj_b = env.sim.get_asset(obj_b_cfg.uid) + # type-check both are RigidObject; else log_error + 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, + ) +``` + +### `remove_rigid_constraint` (detach) + +```python +def remove_rigid_constraint(env, env_ids, name: str) -> None: + """Remove the named constraint for the given env_ids. Idempotent.""" + env.sim.remove_rigid_constraint(name, env_ids=env_ids) +``` + +### Registration & triggering + +```python +@configclass +class MyTaskEventsCfg: + attach_objects: EventCfg = EventCfg( + func=create_rigid_constraint, mode="attach", + params={ + "obj_a_cfg": SceneEntityCfg(uid="cube"), + "obj_b_cfg": SceneEntityCfg(uid="block"), + "name": "cube_block_weld", + }, + ) + detach_objects: EventCfg = EventCfg( + func=remove_rigid_constraint, mode="detach", + params={"name": "cube_block_weld"}, + ) +``` + +```python +self.event_manager.apply(mode="attach", env_ids=gripping_env_ids) +self.event_manager.apply(mode="detach", env_ids=released_env_ids) +``` + +### Decisions + +1. **Thin adapter** — functor does resolution + delegation only; no dexsim calls, no state. +2. **`env_ids` threading** — forwarded to `env.sim.create_rigid_constraint(..., env_ids=...)`. +3. **Custom modes** (`"attach"`/`"detach"`) — task-driven; supported by `EventManager` (arbitrary mode string, task wires `apply`). +4. **Detach takes only `name` + `env_ids`** — looked up by name in the sim registry. + +## 6. Data flow + +### Create (attach) + +``` +task → event_manager.apply(mode="attach", env_ids=gripping_env_ids) + → EventManager iterates "attach" functors → func(env, env_ids, **params) + → create_rigid_constraint(env, env_ids, obj_a_cfg, obj_b_cfg, name, frames) + obj_a = env.sim.get_asset(obj_a_cfg.uid); obj_b = env.sim.get_asset(obj_b_cfg.uid) + build RigidConstraintCfg(...) + → env.sim.create_rigid_constraint(cfg, env_ids) + resolve obj_a, obj_b from self._rigid_objects (raise if missing) + target_env_ids = env_ids or range(num_envs) + frames = broadcast(cfg.local_frame_a) # None→I, 4x4→repeat, (N,4,4)→index + for i in target_env_ids: + arena = self.get_env(i) + name_i = cfg.name if num_envs==1 else f"{cfg.name}_{i}" + handle = arena.create_fixed_constraint(name_i, obj_a[i], obj_b[i], fa, fb) + if handle is None: log_error(arena i) + self._constraints[cfg.name] = RigidConstraint(...) + → physics: obj_a[i] and obj_b[i] welded in arena[i] +``` + +### Remove (detach) + +``` +task → event_manager.apply(mode="detach", env_ids=released_env_ids) + → remove_rigid_constraint(env, env_ids, name) + → env.sim.remove_rigid_constraint(name, env_ids) + constraint = self._constraints.pop(name, None) + if None: log_warning; return False + constraint.destroy(env_ids) # arena.remove_constraint(name_i) per env + → physics: obj_a[i] and obj_b[i] no longer welded +``` + +### Per-env selectivity — index alignment + +`constraint_handles` is a list of length `num_envs`, `None` wherever inactive: + +``` +num_envs = 4, attach on env_ids=[0, 2] +constraint_handles = [handle_0, None, handle_2, None] +``` + +- `get_relative_transform(env_ids=[2])` reads `constraint_handles[2]`. +- `remove(env_ids=[0])` clears index 0 only; index 2 stays attached. +- Partial remove: `destroy(env_ids=[0])` → `arena.remove_constraint("name_0")` + `constraint_handles[0] = None`. When all handles become `None`, the wrapper is dropped from `self._constraints` (base name freed). + +v1 lifecycle is `create` → `remove`. Re-attaching after a partial remove requires removing the whole constraint by name and creating it again (consistent with `add_rigid_object`'s "already exists" semantics — a duplicate base name on `create` is an error, even if only some envs are currently active). + +At create time the `constraint_handles` list is pre-sized to `num_envs` filled with `None`, then only the `target_env_ids` entries are populated, so arena-index alignment always holds. + +### Local-frame broadcasting (once, at create time) + +| Input | Normalized per-env | +|-------|--------------------| +| `None` | `np.eye(4)` for all envs → weld at current relative pose | +| `(4, 4)` | same matrix for all envs | +| `(N, 4, 4)` | `frames[i]` for env `i` (requires N == num_envs) | + +`local_frame_b` mirrors this independently. + +### Interaction with reset + +Constraints are **not** auto-reset by `reset_objects_state` (constraints aren't bodies). Default task policy: register a `reset`-mode `remove_rigid_constraint` (or call `sim.remove_rigid_constraint(name)` in the task's `_reset`) so stale constraints don't leak across episodes. The sim layer does not silently create/destroy on reset. + +## 7. Error handling + +| Condition | Layer | Behavior | +|-----------|-------|----------| +| Either RigidObject uid missing | sim | `log_error` (raises) | +| Duplicate base name in `_constraints` | sim | `log_error` | +| `create_fixed_constraint` returns `None` | sim | `log_error` with arena index | +| `(N,4,4)` frame N ≠ num_envs | sim | `log_error` | +| `remove` on unknown name | sim | `log_warning`, return `False` | +| `remove` with env_ids subset | sim | clear only those handles | +| Entity in functor not a `RigidObject` | functor | `log_error` | +| Remove a constraint already removed (handle `None`) | sim | no-op, success | +| Static + dynamic body combo | sim | allowed (weld-to-environment) | + +## 8. Testing + +**`tests/sim/objects/test_rigid_constraint.py`** (sim layer, mocks — `MockSim` exposing `create_fixed_constraint`/`remove_constraint` returning mock handles; `add_rigid_object` returning mock `RigidObject`s with `_entities`): + +- `test_create_resolves_both_objects` +- `test_create_missing_object_raises` +- `test_create_subset_env_ids` — handles at 0,2 only; `None` elsewhere; alignment holds +- `test_local_frame_broadcasting` — None→identity, 4×4→repeat, (N,4,4)→indexed, bad N→error +- `test_remove_by_name_clears_handles` +- `test_remove_unknown_name_warns_false` +- `test_partial_remove_keeps_others` +- `test_all_removed_drops_from_registry` +- `test_get_relative_transform_skips_none_handles` + +**`tests/gym/envs/managers/test_event_rigid_constraint.py`** (functor layer, mocks — `MockEnv` with `MockSim`, spied `create/remove_rigid_constraint`): + +- `test_create_functor_delegates_to_sim` — params forwarded, `env_ids` threaded +- `test_create_functor_rejects_non_rigid_object` — `Articulation` uid → error +- `test_remove_functor_delegates` — forwards name + env_ids +- `test_custom_mode_apply_invokes_functor` — `EventCfg(mode="attach")` → `apply("attach", env_ids)` calls the spy once with those env_ids + +**`tests/sim/test_rigid_constraint_integration.py`** (real-sim smoke, `@pytest.mark.gpu` / skipped without display): + +- Mirror the dexsim `test_constraint.py` contract at the EmbodiChain layer: attach two dynamic cubes, step, assert relative transform stays constant; detach, step, assert they separate. + +## 9. File layout + +``` +embodichain/lab/sim/objects/constraint.py # RigidConstraint +embodichain/lab/sim/objects/__init__.py # export RigidConstraint +embodichain/lab/sim/cfg.py # + RigidConstraintCfg +embodichain/lab/sim/sim_manager.py # + create/remove/get_rigid_constraint, + # _constraints registry, asset_uids + + # _deferred_destroy wiring +embodichain/lab/gym/envs/managers/events.py # + create/remove_rigid_constraint functors + # + __all__ +tests/sim/objects/test_rigid_constraint.py # sim-layer unit (mocks) +tests/gym/envs/managers/test_event_rigid_constraint.py # functor unit (mocks) +tests/sim/test_rigid_constraint_integration.py # real-sim smoke (gpu-marked) +``` + +## 10. Out of scope / extension points + +- Typed constraints (prismatic/revolute/spherical/d6) → `RigidConstraintCfg.constraint_type` reserved; typed factories land later without public API change. +- Per-constraint physics tuning (limits, drive, motion) → with typed constraints. +- Articulation constraints → `RigidObject`↔`RigidObject` only for v1. +- Auto-detach-on-reset → task policy, not sim policy. +- A dedicated `/add-...` skill → not needed; functors follow existing `add-functor` conventions. From c9a21304892e4c5d1298c924a34b34f3a0d17b3a Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:05:37 +0000 Subject: [PATCH 02/16] Add rigid constraint implementation plan Co-Authored-By: Claude --- .../plans/2026-06-29-rigid-constraint.md | 1691 +++++++++++++++++ 1 file changed, 1691 insertions(+) create mode 100644 docs/superpowers/plans/2026-06-29-rigid-constraint.md diff --git a/docs/superpowers/plans/2026-06-29-rigid-constraint.md b/docs/superpowers/plans/2026-06-29-rigid-constraint.md new file mode 100644 index 00000000..d7126bd8 --- /dev/null +++ b/docs/superpowers/plans/2026-06-29-rigid-constraint.md @@ -0,0 +1,1691 @@ +# Rigid Object Constraint Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add the ability to attach two `RigidObject`s via a fixed physics constraint and remove it, exposed both as a standalone `SimulationManager` API and as on-demand event functors triggered from a task environment. + +**Architecture:** A sim-layer `RigidConstraint` batch wrapper mirrors `RigidObject`'s per-arena pattern, holding one dexsim `FixedConstraint` handle per arena (with `None` where inactive, so arena-index == list-index). `SimulationManager` owns the constraint registry and all dexsim calls. Two function-style event functors in `events.py` (`create_rigid_constraint`, `remove_rigid_constraint`) resolve `SceneEntityCfg` → `RigidObject` and delegate to the sim API, triggered via custom event modes (`"attach"`/`"detach"`). + +**Tech Stack:** Python 3.11, `embodichain` package, `dexsim` (Warp/PhysX), `torch`, `numpy`, `pytest`, `black==26.3.1`. + +## Global Constraints + +- Every source file begins with the Apache 2.0 copyright header (see `CLAUDE.md`). +- `from __future__ import annotations` at the top of every new file. +- Use `TYPE_CHECKING` guard for `EmbodiedEnv` / `SimulationManager` imports to avoid circular imports. +- Prefer `A | B` over `Union[A, B]`. +- `@configclass` for all config objects; `MISSING` for required fields. +- `logger.log_error(msg)` raises `RuntimeError` by default — error paths call it and let it raise (matches `add_rigid_object`). +- Formatter: `black==26.3.1`. Run `black ` before each commit. +- Package name is `embodichain` (lowercase). +- Spec: `docs/superpowers/specs/2026-06-29-rigid-constraint-design.md`. + +--- + +## File Structure + +``` +embodichain/lab/sim/cfg.py # MODIFY: add RigidConstraintCfg +embodichain/lab/sim/objects/constraint.py # CREATE: RigidConstraint wrapper +embodichain/lab/sim/objects/__init__.py # MODIFY: export RigidConstraint +embodichain/lab/sim/sim_manager.py # MODIFY: +_constraints registry, + # create/remove/get_rigid_constraint, + # asset_uids + _deferred_destroy wiring +embodichain/lab/gym/envs/managers/events.py # MODIFY: +2 functors (no __all__ exists) +tests/sim/objects/test_rigid_constraint.py # CREATE: sim-layer unit tests (mocks) +tests/gym/envs/managers/test_event_rigid_constraint.py # CREATE: functor unit tests (mocks) +tests/sim/test_rigid_constraint_integration.py # CREATE: real-sim smoke (gpu-marked) +``` + +`RigidConstraintCfg` lives in `cfg.py` (where all sim cfgs live). `RigidConstraint` lives in `objects/constraint.py`. The functors live in `events.py` (no `__all__` in that file — just add the functions). + +--- + +### Task 1: `RigidConstraintCfg` in `cfg.py` + +**Files:** +- Modify: `embodichain/lab/sim/cfg.py` (append a new `@configclass` near `RigidObjectGroupCfg` ~line 900) +- Test: `tests/sim/objects/test_rigid_constraint.py` (create file, first test) + +**Interfaces:** +- Produces: `RigidConstraintCfg` dataclass with fields `name: str = MISSING`, `rigid_object_a_uid: str = MISSING`, `rigid_object_b_uid: str = MISSING`, `local_frame_a: np.ndarray | None = None`, `local_frame_b: np.ndarray | None = None`, `constraint_type: Literal["fixed"] = "fixed"`. + +- [ ] **Step 1: Write the failing test** + +Create `tests/sim/objects/test_rigid_constraint.py`: + +```python +# ---------------------------------------------------------------------------- +# 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 + +from dataclasses import MISSING + +from embodichain.lab.sim.cfg import RigidConstraintCfg + + +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) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` +Expected: FAIL with `ImportError: cannot import name 'RigidConstraintCfg'` + +- [ ] **Step 3: Write minimal implementation** + +Append to `embodichain/lab/sim/cfg.py` (right after the `RigidObjectGroupCfg` class ends, ~line 1000, before the next `@configclass`). Note `Literal` and `Optional` are already imported at the top of `cfg.py` (lines 22): `from typing import Sequence, Union, Dict, Literal, List, Any, Optional`. + +```python +@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`` attaches at the objects' current relative pose (identity). + 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: As :attr:`local_frame_a`, for object B. 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 (current relative pose).""" + + local_frame_b: np.ndarray | None = None + """Local joint frame on object B. None -> identity (current relative pose).""" + + constraint_type: Literal["fixed"] = "fixed" + """Constraint type. Only ``"fixed"`` is supported in v1.""" +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` +Expected: PASS (3 tests) + +- [ ] **Step 5: Commit** + +```bash +black embodichain/lab/sim/cfg.py tests/sim/objects/test_rigid_constraint.py +git add embodichain/lab/sim/cfg.py tests/sim/objects/test_rigid_constraint.py +git commit -m "feat(sim): add RigidConstraintCfg" +``` + +--- + +### Task 2: `RigidConstraint` wrapper in `objects/constraint.py` + +**Files:** +- Create: `embodichain/lab/sim/objects/constraint.py` +- Modify: `embodichain/lab/sim/objects/__init__.py` (export `RigidConstraint`) +- Test: `tests/sim/objects/test_rigid_constraint.py` (append) + +**Interfaces:** +- Consumes: `RigidConstraintCfg` (Task 1), `RigidObject` (from `embodichain.lab.sim.objects`). +- Produces: `RigidConstraint` class with: + - `__init__(cfg: RigidConstraintCfg, constraint_handles: list, rigid_object_a: RigidObject, rigid_object_b: RigidObject, device: torch.device)` + - `@property num_envs -> int` + - `get_relative_transform(env_ids=None) -> list[np.ndarray]` + - `get_local_pose(actor_index: int, env_ids=None) -> list[np.ndarray]` + - `get_name(env_id: int) -> str` (returns the base name for single env, `f"{base}_{env_id}"` for multi) + - `is_valid(env_ids=None) -> list[bool]` + - `destroy(self, env_ids=None, arena_resolver=None) -> None` — calls `arena.remove_constraint(name_i)` per env; `arena_resolver(i)` returns the arena for env `i`. When all handles become None, the caller (SimulationManager) drops the wrapper. + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/sim/objects/test_rigid_constraint.py`: + +```python +import torch +from unittest.mock import MagicMock + +from embodichain.lab.sim.objects.constraint import RigidConstraint + + +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) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` +Expected: FAIL with `ImportError: cannot import name 'RigidConstraint'` + +- [ ] **Step 3: Write minimal implementation** + +Create `embodichain/lab/sim/objects/constraint.py`: + +```python +# ---------------------------------------------------------------------------- +# 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 + + +@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. + rigid_object_b: The second RigidObject. + device: The torch device. + """ + + cfg: RigidConstraintCfg + constraint_handles: list[Any] = field(default_factory=list) + rigid_object_a: RigidObject = None + rigid_object_b: RigidObject = 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 +``` + +Then modify `embodichain/lab/sim/objects/__init__.py` — add the import and keep it exported. Add after the `from .gizmo import Gizmo` line (line 29): + +```python +from .constraint import RigidConstraint +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` +Expected: PASS (all tests, including the new ones) + +- [ ] **Step 5: Commit** + +```bash +black embodichain/lab/sim/objects/constraint.py embodichain/lab/sim/objects/__init__.py tests/sim/objects/test_rigid_constraint.py +git add embodichain/lab/sim/objects/constraint.py embodichain/lab/sim/objects/__init__.py tests/sim/objects/test_rigid_constraint.py +git commit -m "feat(sim): add RigidConstraint wrapper" +``` + +--- + +### Task 3: `SimulationManager` constraint registry + create API + +**Files:** +- Modify: `embodichain/lab/sim/sim_manager.py` (add `_constraints` registry in `__init__`, `create_rigid_constraint`, plus helpers) +- Test: `tests/sim/objects/test_rigid_constraint.py` (append sim-layer tests using a mock sim) + +**Interfaces:** +- Consumes: `RigidConstraintCfg` (Task 1), `RigidConstraint` (Task 2). +- Produces: `SimulationManager.create_rigid_constraint(cfg, env_ids=None) -> RigidConstraint`. The method resolves objects from `self._rigid_objects`, broadcasts local frames, and for each target env calls `self.get_env(i).create_fixed_constraint(name_i, obj_a._entities[i], obj_b._entities[i], frame_a, frame_b)`, stores a `RigidConstraint` in `self._constraints[cfg.name]`. + +**Frame broadcasting rules** (a helper `_broadcast_frame`): `None` -> `np.eye(4)` per env; `(4,4)` -> same matrix per env; `(N,4,4)` -> index `i` (requires `N == num_envs`, else `log_error`). + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/sim/objects/test_rigid_constraint.py`: + +```python +from embodichain.lab.sim.sim_manager import SimulationManager +from embodichain.utils import configclass # noqa: F401 (ensures import works) + + +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 + _broadcast_frame = staticmethod(SimulationManager._broadcast_frame) + + +def _register_object(sim, uid, num_envs): + obj = MagicMock() + obj.uid = uid + obj.num_instances = num_envs + obj._entities = [MagicMock(name=f"{uid}_{i}") for i in range(num_envs)] + 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[i][0] == f"weld_{i}" + assert arena.created[i][1] is sim._rigid_objects["cube"]._entities[i] + assert arena.created[i][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") +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py -v -k "create or broadcast"` +Expected: FAIL with `AttributeError: 'SimulationManager' has no attribute 'create_rigid_constraint'` (and `_broadcast_frame`). + +- [ ] **Step 3: Write minimal implementation** + +In `embodichain/lab/sim/sim_manager.py`: + +3a. Add to the imports block near the top (after the existing `from embodichain.lab.sim.cfg import ...` block, ~line 89). Add `RigidConstraintCfg` to that import: + +```python +from embodichain.lab.sim.cfg import ( + RenderCfg, + PhysicsCfg, + MarkerCfg, + GPUMemoryCfg, + WindowRecordCfg, + LightCfg, + RigidObjectCfg, + SoftObjectCfg, + ClothObjectCfg, + RigidObjectGroupCfg, + ArticulationCfg, + RobotCfg, + RigidConstraintCfg, +) +``` + +3b. Add `RigidConstraint` to the objects import (line ~59-67): + +```python +from embodichain.lab.sim.objects import ( + RigidObject, + RigidObjectGroup, + SoftObject, + ClothObject, + Articulation, + Robot, + Light, + RigidConstraint, +) +``` + +3c. In `__init__`, add the registry next to `self._rigid_objects` (line ~285): + +```python + self._rigid_objects: Dict[str, RigidObject] = dict() + self._constraints: Dict[str, RigidConstraint] = dict() +``` + +3d. Add the methods. Place them right after `get_rigid_object_uid_list` (line ~1005, before `add_rigid_object_group`). + +```python + @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)." + ) + + def create_rigid_constraint( + self, + cfg: RigidConstraintCfg, + env_ids: Sequence[int] | 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 + identity (attach at the objects' current relative pose). + + Args: + cfg: The constraint configuration. + env_ids: Target environment indices. 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 + if env_ids is None: + target_env_ids = list(range(num_envs)) + else: + target_env_ids = list(env_ids) + + # broadcast local frames + frames_a = self._broadcast_frame( + cfg.local_frame_a, num_envs, target_env_ids, cfg.name + ) + 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 +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` +Expected: PASS (all tests so far) + +- [ ] **Step 5: Commit** + +```bash +black embodichain/lab/sim/sim_manager.py tests/sim/objects/test_rigid_constraint.py +git add embodichain/lab/sim/sim_manager.py tests/sim/objects/test_rigid_constraint.py +git commit -m "feat(sim): add SimulationManager.create_rigid_constraint" +``` + +--- + +### Task 4: `remove_rigid_constraint` + `get_rigid_constraint` + registry wiring + +**Files:** +- Modify: `embodichain/lab/sim/sim_manager.py` (add remove/get methods; wire `asset_uids` + `_deferred_destroy`) +- Test: `tests/sim/objects/test_rigid_constraint.py` (append) + +**Interfaces:** +- Consumes: `RigidConstraint.destroy` (Task 2), `self._constraints` (Task 3). +- Produces: + - `remove_rigid_constraint(name, env_ids=None) -> bool` — pops (or partially clears) the constraint; calls `constraint.destroy(env_ids, arena_resolver=self.get_env)`. When all handles are None, drops from registry. Returns True if the constraint existed (or still partially exists after subset remove), False if name unknown. + - `get_rigid_constraint(name) -> RigidConstraint | None` + - `get_rigid_constraint_uid_list() -> list[str]` + - `asset_uids` extended to include constraint names. + - `_deferred_destroy` severs `_constraints`. + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/sim/objects/test_rigid_constraint.py`: + +```python +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") + constraint = 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 +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py -v -k "remove or get_rigid"` +Expected: FAIL with `AttributeError: ... has no attribute 'remove_rigid_constraint'`. + +- [ ] **Step 3: Write minimal implementation** + +In `embodichain/lab/sim/sim_manager.py`, add after `create_rigid_constraint` (from Task 3): + +```python + def remove_rigid_constraint( + self, + name: str, + env_ids: Sequence[int] | 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. 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 + + constraint.destroy(env_ids=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 of constraint names. + """ + return list(self._constraints.keys()) +``` + +Then wire `asset_uids` (line ~421). Add constraints to the returned list. In `asset_uids`: + +```python + @property + def asset_uids(self) -> List[str]: + """Get all assets uid in the simulation. + + The assets include lights, sensors, robots, rigid objects and articulations. + + Returns: + List[str]: list of all assets uid. + """ + uid_list = ["default_plane"] + uid_list.extend(list(self._lights.keys())) + uid_list.extend(list(self._sensors.keys())) + uid_list.extend(list(self._robots.keys())) + uid_list.extend(list(self._rigid_objects.keys())) + uid_list.extend(list(self._rigid_object_groups.keys())) + 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 +``` + +Then wire `_deferred_destroy`. In `_deferred_destroy` (~line 2271, the `_sever_wrapper_refs` calls), add: + +```python + _sever_wrapper_refs("_constraints") +``` + +after the `_sever_wrapper_refs("_rigid_object_groups")` line. Also clear `_constraints` in the explicit `clear()` block — find the line `self._rigid_object_groups` clear region and add: + +```python + self._constraints.clear() +``` + +near the other `.clear()` calls at the end of `_deferred_destroy` (after `self._arenas.clear()`). + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` +Expected: PASS (all sim-layer tests) + +- [ ] **Step 5: Commit** + +```bash +black embodichain/lab/sim/sim_manager.py tests/sim/objects/test_rigid_constraint.py +git add embodichain/lab/sim/sim_manager.py tests/sim/objects/test_rigid_constraint.py +git commit -m "feat(sim): add remove/get_rigid_constraint and registry wiring" +``` + +--- + +### Task 5: Event functors in `events.py` + +**Files:** +- Modify: `embodichain/lab/gym/envs/managers/events.py` (add `create_rigid_constraint` + `remove_rigid_constraint` functions) +- Test: `tests/gym/envs/managers/test_event_rigid_constraint.py` (create) + +**Interfaces:** +- Consumes: `SceneEntityCfg` (from `.cfg`), `RigidConstraintCfg` (from `embodichain.lab.sim.cfg`), `RigidObject` (for isinstance check), `logger`. +- Produces: two module-level functions in `events.py`: + - `create_rigid_constraint(env, env_ids, obj_a_cfg: SceneEntityCfg, obj_b_cfg: SceneEntityCfg, name: str, local_frame_a=None, local_frame_b=None) -> None` + - `remove_rigid_constraint(env, env_ids, name: str) -> None` + +- [ ] **Step 1: Write the failing tests** + +Create `tests/gym/envs/managers/test_event_rigid_constraint.py`: + +```python +# ---------------------------------------------------------------------------- +# 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.lab.gym.envs.managers.cfg import SceneEntityCfg +from embodichain.lab.gym.envs.managers.events import ( + create_rigid_constraint, + remove_rigid_constraint, +) +from embodichain.lab.sim.objects.rigid_object import RigidObject + + +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.""" + env = MagicMock() + env.device = torch.device("cpu") + env.num_envs = 4 + + obj_a = MagicMock() + obj_b = MagicMock() + env.sim.get_asset.side_effect = lambda uid: {"cube": obj_a, "block": obj_b}[uid] + + # control the isinstance check by patching RigidObject for the test + env._obj_a_is_rigid = obj_a_is_rigid + env._obj_b_is_rigid = obj_b_is_rigid + 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).""" + env, obj_a, obj_b = _make_env() + # patch isinstance check to return False for obj_a + monkeypatch.setattr( + "embodichain.lab.gym.envs.managers.events.RigidObject", MagicMock(__instancecheck__=lambda self, o: False) + ) + 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) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `pytest tests/gym/envs/managers/test_event_rigid_constraint.py -v` +Expected: FAIL with `ImportError: cannot import name 'create_rigid_constraint'` + +- [ ] **Step 3: Write minimal implementation** + +In `embodichain/lab/gym/envs/managers/events.py`, add the imports near the top (the file already imports `RigidObject` at line 28, `logger`, `np`-equivalent math helpers, and `SceneEntityCfg` at line 35). Add to the existing `from embodichain.lab.sim.objects import (...)` block a new import is not needed (RigidObject already there). Add `RigidConstraintCfg` import. Find the `from embodichain.lab.sim.cfg import RigidObjectCfg, ArticulationCfg` line (line 33) and extend it: + +```python +from embodichain.lab.sim.cfg import RigidObjectCfg, ArticulationCfg, RigidConstraintCfg +``` + +Then add `numpy` import if not present (the file uses math helpers but check). At the top, after `import random` (line 21), ensure `import numpy as np` exists. If not present, add it. + +Then append the two functors at the end of the file (after `set_detached_uids_for_env_reset`): + +```python +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 attaches at the + objects' current relative pose. Accepts (4,4) or (N,4,4). + local_frame_b: Local joint frame on object B. None -> identity. + + 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) +``` + +Note: `EmbodiedEnv` is already imported under `TYPE_CHECKING` at the top of `events.py` (line 50). `torch` is already imported (line 19). + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `pytest tests/gym/envs/managers/test_event_rigid_constraint.py -v` +Expected: PASS (5 tests) + +- [ ] **Step 5: Commit** + +```bash +black embodichain/lab/gym/envs/managers/events.py tests/gym/envs/managers/test_event_rigid_constraint.py +git add embodichain/lab/gym/envs/managers/events.py tests/gym/envs/managers/test_event_rigid_constraint.py +git commit -m "feat(env): add rigid-constraint event functors" +``` + +--- + +### Task 6: EventManager custom-mode wiring test + +**Files:** +- Test: `tests/gym/envs/managers/test_event_rigid_constraint.py` (append) + +**Interfaces:** +- Consumes: `EventManager`, `EventCfg`, the two functors (Task 5), `ManagerBase`. +- Produces: a test proving `EventManager.apply(mode="attach", env_ids)` invokes a registered custom-mode functor once with those env_ids (no env subclassing needed — the manager supports arbitrary mode strings). + +- [ ] **Step 1: Write the failing test** + +Append to `tests/gym/envs/managers/test_event_rigid_constraint.py`: + +```python +from embodichain.lab.gym.envs.managers.event_manager import EventManager +from embodichain.lab.gym.envs.managers.cfg import EventCfg +from embodichain.utils import configclass + + +@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() + 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 +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `pytest tests/gym/envs/managers/test_event_rigid_constraint.py::test_custom_mode_apply_invokes_functor_with_env_ids -v` +Expected: FAIL — likely an import or attribute error if `EventManager` construction with the mock env fails. If the mock env is insufficient, fix the mock (the test asserts behavior, so adjust the mock to satisfy EventManager's needs). Inspect the failure and patch the mock accordingly (e.g. add `env.cfg = None` if needed). + +- [ ] **Step 3: (No new implementation needed)** + +`EventManager` already supports arbitrary mode strings (see `event_manager.py` `_prepare_functors` — it registers any `functor_cfg.mode` into `_mode_functor_names`). If the test fails on a missing mock attribute, fix the test's mock rather than changing `EventManager`. Document what was needed. + +- [ ] **Step 4: Run test to verify it passes** + +Run: `pytest tests/gym/envs/managers/test_event_rigid_constraint.py -v` +Expected: PASS (6 tests) + +- [ ] **Step 5: Commit** + +```bash +black tests/gym/envs/managers/test_event_rigid_constraint.py +git add tests/gym/envs/managers/test_event_rigid_constraint.py +git commit -m "test(env): custom-mode apply invokes rigid-constraint functor" +``` + +--- + +### Task 7: Real-sim integration smoke test + +**Files:** +- Create: `tests/sim/test_rigid_constraint_integration.py` + +**Interfaces:** +- Consumes: `SimulationManager`, `RigidObjectCfg`, `RigidConstraintCfg`, real dexsim. Mirrors the contract in `dexsim/python/test/engine/test_constraint.py` at the EmbodiChain layer. +- Produces: a skipped-unless-gpu test that attaches two dynamic cubes, steps, asserts the relative transform stays constant, detaches, steps, asserts separation. + +- [ ] **Step 1: Write the test** + +Create `tests/sim/test_rigid_constraint_integration.py`: + +```python +# ---------------------------------------------------------------------------- +# 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. + +Skipped unless a GPU/display is available. Mirrors the dexsim +test_constraint.py contract at the EmbodiChain SimulationManager layer. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from embodichain.lab.sim.sim_manager import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import ( + RigidObjectCfg, + RigidConstraintCfg, + RigidBodyAttributesCfg, +) +from embodichain.lab.sim.shapes import MeshCfg +from dexsim.types import ActorType, RigidBodyShape + + +def _can_run_gpu_sim() -> bool: + try: + return torch.cuda.is_available() + except Exception: + return False + + +pytestmark = pytest.mark.skipif( + not _can_run_gpu_sim(), reason="GPU simulation required for constraint integration test" +) + + +def _make_sim(num_envs: int = 1) -> SimulationManager: + cfg = SimulationManagerCfg() + cfg.headless = True + cfg.sim_device = "cuda" + cfg.num_envs = num_envs + return SimulationManager(cfg) + + +def test_fixed_constraint_holds_relative_pose(): + """Two welded cubes keep their relative transform under physics; detach lets them separate.""" + sim = _make_sim(num_envs=1) + try: + attrs = RigidBodyAttributesCfg() + attrs.mass = 0.2 + cube_cfg = RigidObjectCfg( + uid="cube_a", + body_type="dynamic", + init_pos=[0.0, 0.0, 1.4], + attrs=attrs, + shape=MeshCfg(fpath="..."), # placeholder; see note + ) + block_cfg = RigidObjectCfg( + uid="cube_b", + body_type="dynamic", + init_pos=[0.0, 0.0, 1.2], + attrs=attrs, + shape=MeshCfg(fpath="..."), + ) + cube = sim.add_rigid_object(cube_cfg) + block = sim.add_rigid_object(block_cfg) + + constraint = sim.create_rigid_constraint( + RigidConstraintCfg( + name="weld", + rigid_object_a_uid="cube_a", + rigid_object_b_uid="cube_b", + ) + ) + assert constraint.is_valid() == [True] + + initial_delta_z = ( + block.get_local_pose()[0, 2, 3] - cube.get_local_pose()[0, 2, 3] + ) + + for _ in range(120): + sim.update(step=1) + + rel = constraint.get_relative_transform()[0] + np.testing.assert_allclose(rel[:3, 3], np.zeros(3), atol=2e-2) + delta_z = ( + block.get_local_pose()[0, 2, 3] - cube.get_local_pose()[0, 2, 3] + ) + assert abs(delta_z - initial_delta_z) < 2e-2 + + # detach and confirm they separate + sim.remove_rigid_constraint("weld") + z_before = cube.get_local_pose()[0, 2, 3] + for _ in range(120): + sim.update(step=1) + z_after = cube.get_local_pose()[0, 2, 3] + assert z_after < z_before - 0.05 + finally: + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() +``` + +- [ ] **Step 2: Run test to verify it skips (or passes on GPU)** + +Run: `pytest tests/sim/test_rigid_constraint_integration.py -v` +Expected: SKIPPED (no GPU) on CPU CI. On a GPU machine it should PASS. + +- [ ] **Step 3: (Implementation already complete from Tasks 1-5)** + +This test exercises the full sim path. If `RigidObjectCfg`/`MeshCfg` construction needs a real mesh path, replace the `fpath="..."` placeholders with a real asset path from `embodichain/data/assets` (check `SimResources` for a built-in cube mesh). Document the chosen path in the test. If the test cannot run headless on the CI runner, leave it gpu-marked and skipped — that's acceptable per the spec. + +- [ ] **Step 4: Commit** + +```bash +black tests/sim/test_rigid_constraint_integration.py +git add tests/sim/test_rigid_constraint_integration.py +git commit -m "test(sim): add rigid-constraint real-sim integration smoke test" +``` + +--- + +### Task 8: Run full suite, black, finalize + +**Files:** +- All touched files. + +- [ ] **Step 1: Run the full constraint test suite** + +Run: `pytest tests/sim/objects/test_rigid_constraint.py tests/gym/envs/managers/test_event_rigid_constraint.py tests/sim/test_rigid_constraint_integration.py -v` +Expected: all unit tests PASS; integration test SKIPPED (or PASS on GPU). + +- [ ] **Step 2: Run black on all changed files** + +Run: `black embodichain/lab/sim/cfg.py embodichain/lab/sim/objects/constraint.py embodichain/lab/sim/objects/__init__.py embodichain/lab/sim/sim_manager.py embodichain/lab/gym/envs/managers/events.py tests/sim/objects/test_rigid_constraint.py tests/gym/envs/managers/test_event_rigid_constraint.py tests/sim/test_rigid_constraint_integration.py` +Expected: no changes (or only formatting). + +- [ ] **Step 3: Run the pre-commit check skill** + +Use the `/pre-commit-check` skill (it runs all local CI checks). Fix any violations it reports. + +- [ ] **Step 4: Commit any formatting fixes** + +```bash +git add -A +git commit -m "chore: black formatting + pre-commit fixes" +``` + +- [ ] **Step 5: Final summary** + +The implementation is complete. Summary of what was built: +- `RigidConstraintCfg` (cfg.py) +- `RigidConstraint` wrapper (objects/constraint.py) +- `SimulationManager.create_rigid_constraint` / `remove_rigid_constraint` / `get_rigid_constraint` + `_constraints` registry +- `create_rigid_constraint` / `remove_rigid_constraint` event functors (events.py) +- Unit tests (mocks) + integration smoke test (gpu-marked) + +## Self-Review + +**1. Spec coverage:** +- §3 Architecture (sim layer + functor layer): Tasks 1-5. +- §4 Sim-layer API (RigidConstraint, RigidConstraintCfg, SimulationManager methods): Tasks 1-4. +- §5 Functor layer (two functors, registration/triggering): Tasks 5-6. +- §6 Data flow (create/remove, per-env selectivity, frame broadcasting, reset interaction): Tasks 3-4 + Task 6 (reset interaction is task-policy, documented in spec; no sim code needed). +- §7 Error handling: each error case is a test in Task 3-4. +- §8 Testing: Tasks 3 (sim unit), 5-6 (functor unit), 7 (integration). +- §9 File layout: matches. + +**2. Placeholder scan:** The `MeshCfg(fpath="...")` in Task 7 is flagged as needing a real asset path — that's an integration-test asset-resolution detail, addressed in Task 7 Step 3. No "TBD"/"implement later" in deliverable code. + +**3. Type consistency:** `create_rigid_constraint` / `remove_rigid_constraint` / `get_rigid_constraint` names match across sim layer (Task 3-4) and functor layer (Task 5). `RigidConstraint` field names (`constraint_handles`, `rigid_object_a`, `rigid_object_b`, `device`) match Task 2 init and Task 3 usage. `destroy(env_ids, arena_resolver)` signature matches Task 2 and Task 4's call. `_broadcast_frame` is a `@staticmethod` accessed as `SimulationManager._broadcast_frame` in tests (Task 3) and as `self._broadcast_frame` in `create_rigid_constraint` — both valid. From f231dc42da0c4a444fb0495e32d44e698855109c Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:07:15 +0000 Subject: [PATCH 03/16] feat(sim): add RigidConstraintCfg Co-Authored-By: Claude --- embodichain/lab/sim/cfg.py | 44 +++++++++++++++ tests/sim/objects/test_rigid_constraint.py | 65 ++++++++++++++++++++++ 2 files changed, 109 insertions(+) create mode 100644 tests/sim/objects/test_rigid_constraint.py diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 0b2fe922..ce137e16 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -997,6 +997,50 @@ 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`` attaches at the objects' current relative pose (identity). + 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: As :attr:`local_frame_a`, for object B. 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 (current relative pose).""" + + local_frame_b: np.ndarray | None = None + """Local joint frame on object B. None -> identity (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/tests/sim/objects/test_rigid_constraint.py b/tests/sim/objects/test_rigid_constraint.py new file mode 100644 index 00000000..d94db74d --- /dev/null +++ b/tests/sim/objects/test_rigid_constraint.py @@ -0,0 +1,65 @@ +# ---------------------------------------------------------------------------- +# 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 + +from dataclasses import MISSING + +from embodichain.lab.sim.cfg import RigidConstraintCfg + + +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) From 5221c163e5cd4e9207da8d86ef6ae0ec3bad46e4 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:09:06 +0000 Subject: [PATCH 04/16] feat(sim): add RigidConstraint wrapper Co-Authored-By: Claude --- embodichain/lab/sim/objects/__init__.py | 1 + embodichain/lab/sim/objects/constraint.py | 147 +++++++++++++++++++ tests/sim/objects/test_rigid_constraint.py | 159 +++++++++++++++++++++ 3 files changed, 307 insertions(+) create mode 100644 embodichain/lab/sim/objects/constraint.py 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..c45fa24d --- /dev/null +++ b/embodichain/lab/sim/objects/constraint.py @@ -0,0 +1,147 @@ +# ---------------------------------------------------------------------------- +# 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 + + +@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. + rigid_object_b: The second RigidObject. + device: The torch device. + """ + + cfg: RigidConstraintCfg + constraint_handles: list[Any] = field(default_factory=list) + rigid_object_a: RigidObject = None + rigid_object_b: RigidObject = 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/tests/sim/objects/test_rigid_constraint.py b/tests/sim/objects/test_rigid_constraint.py index d94db74d..bc335ee8 100644 --- a/tests/sim/objects/test_rigid_constraint.py +++ b/tests/sim/objects/test_rigid_constraint.py @@ -20,10 +20,13 @@ 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 def test_rigid_constraint_cfg_defaults(): @@ -63,3 +66,159 @@ def test_rigid_constraint_cfg_accepts_frames(): 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) From 3109c16d12f242a2173d44f6641a39d8945905ee Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:12:18 +0000 Subject: [PATCH 05/16] feat(sim): add SimulationManager.create_rigid_constraint Co-Authored-By: Claude --- embodichain/lab/sim/sim_manager.py | 146 +++++++++++++++ tests/sim/objects/test_rigid_constraint.py | 196 +++++++++++++++++++++ 2 files changed, 342 insertions(+) diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index e8ac2916..b95c07ca 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,6 +87,7 @@ RigidObjectGroupCfg, ArticulationCfg, RobotCfg, + RigidConstraintCfg, ) from embodichain.lab.sim import VisualMaterial, VisualMaterialCfg from embodichain.utils import configclass, logger @@ -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() @@ -1004,6 +1007,149 @@ 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)." + ) + + def create_rigid_constraint( + self, + cfg: RigidConstraintCfg, + env_ids: Sequence[int] | 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 + identity (attach at the objects' current relative pose). + + Args: + cfg: The constraint configuration. + env_ids: Target environment indices. 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 + if env_ids is None: + target_env_ids = list(range(num_envs)) + else: + target_env_ids = list(env_ids) + + # broadcast local frames + frames_a = self._broadcast_frame( + cfg.local_frame_a, num_envs, target_env_ids, cfg.name + ) + 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 diff --git a/tests/sim/objects/test_rigid_constraint.py b/tests/sim/objects/test_rigid_constraint.py index bc335ee8..516b616e 100644 --- a/tests/sim/objects/test_rigid_constraint.py +++ b/tests/sim/objects/test_rigid_constraint.py @@ -222,3 +222,199 @@ def test_rigid_constraint_destroy_all_returns_all_cleared(): 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) + + +from embodichain.lab.sim.sim_manager import SimulationManager + + +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 + _broadcast_frame = staticmethod(SimulationManager._broadcast_frame) + + +def _register_object(sim, uid, num_envs): + obj = MagicMock() + obj.uid = uid + obj.num_instances = num_envs + obj._entities = [MagicMock(name=f"{uid}_{i}") for i in range(num_envs)] + 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") From e17c9a039cce71ccb00c5ba5de7726dbfa062c24 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:15:45 +0000 Subject: [PATCH 06/16] feat(sim): add remove/get_rigid_constraint and registry wiring Co-Authored-By: Claude --- embodichain/lab/sim/sim_manager.py | 56 +++++++++++++++ tests/sim/objects/test_rigid_constraint.py | 80 ++++++++++++++++++++++ 2 files changed, 136 insertions(+) diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index b95c07ca..07c7cbad 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -438,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( @@ -1166,6 +1167,59 @@ 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] | 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. 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 + + constraint.destroy(env_ids=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. @@ -2417,6 +2471,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") @@ -2439,6 +2494,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/tests/sim/objects/test_rigid_constraint.py b/tests/sim/objects/test_rigid_constraint.py index 516b616e..1fe8e9ca 100644 --- a/tests/sim/objects/test_rigid_constraint.py +++ b/tests/sim/objects/test_rigid_constraint.py @@ -279,6 +279,9 @@ def get_env(self, arena_index=-1): # 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) @@ -418,3 +421,80 @@ def test_broadcast_frame_N4x4_indexes(): 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_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 From b509ce46ca322aa2019af574706a777810c58230 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:18:17 +0000 Subject: [PATCH 07/16] feat(env): add rigid-constraint event functors Co-Authored-By: Claude --- embodichain/lab/gym/envs/managers/events.py | 69 ++++++++- .../managers/test_event_rigid_constraint.py | 143 ++++++++++++++++++ 2 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 tests/gym/envs/managers/test_event_rigid_constraint.py diff --git a/embodichain/lab/gym/envs/managers/events.py b/embodichain/lab/gym/envs/managers/events.py index e99f1458..4122ddca 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,68 @@ 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 attaches at the + objects' current relative pose. Accepts (4,4) or (N,4,4). + local_frame_b: Local joint frame on object B. None -> identity. + + 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/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..c08bc8a9 --- /dev/null +++ b/tests/gym/envs/managers/test_event_rigid_constraint.py @@ -0,0 +1,143 @@ +# ---------------------------------------------------------------------------- +# 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.lab.gym.envs.managers.cfg import SceneEntityCfg +from embodichain.lab.gym.envs.managers.events import ( + create_rigid_constraint, + remove_rigid_constraint, +) +from embodichain.lab.sim.objects.rigid_object import RigidObject + + +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) From faa7ac5720acd36874d2792818bfdc751cce0087 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:19:38 +0000 Subject: [PATCH 08/16] test(env): custom-mode apply invokes rigid-constraint functor Co-Authored-By: Claude --- .../managers/test_event_rigid_constraint.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/gym/envs/managers/test_event_rigid_constraint.py b/tests/gym/envs/managers/test_event_rigid_constraint.py index c08bc8a9..b80ac704 100644 --- a/tests/gym/envs/managers/test_event_rigid_constraint.py +++ b/tests/gym/envs/managers/test_event_rigid_constraint.py @@ -141,3 +141,50 @@ def test_remove_functor_none_env_ids(): env, _, _ = _make_env() remove_rigid_constraint(env, None, name="weld") env.sim.remove_rigid_constraint.assert_called_once_with("weld", env_ids=None) + + +from embodichain.lab.gym.envs.managers.event_manager import EventManager +from embodichain.lab.gym.envs.managers.cfg import EventCfg +from embodichain.utils import configclass + + +@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 From c54b04c872b6d4f63f6a217fbdf068e6f12696f7 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:27:23 +0000 Subject: [PATCH 09/16] test(sim): add rigid-constraint real-sim integration smoke test Skips gracefully when the test asset is absent or CUDA is unavailable. On a machine with the asset + GPU, asserts the fixed constraint holds the relative transform under physics and that detaching lets objects separate. Co-Authored-By: Claude --- .../sim/test_rigid_constraint_integration.py | 142 ++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 tests/sim/test_rigid_constraint_integration.py diff --git a/tests/sim/test_rigid_constraint_integration.py b/tests/sim/test_rigid_constraint_integration.py new file mode 100644 index 00000000..2005f5f5 --- /dev/null +++ b/tests/sim/test_rigid_constraint_integration.py @@ -0,0 +1,142 @@ +# ---------------------------------------------------------------------------- +# 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 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, welded at identity frames. + 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 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 they separate (both fall independently). + self.sim.remove_rigid_constraint("weld") + assert self.sim.get_rigid_constraint("weld") is None + z_before = float(self.duck_a.get_local_pose(to_matrix=True)[0, 2, 3]) + for _ in range(120): + self.sim.update(step=1) + z_after = float(self.duck_a.get_local_pose(to_matrix=True)[0, 2, 3]) + assert z_after < z_before - 0.05 + + +class TestRigidConstraintCPU(BaseRigidConstraintTest): + def setup_method(self) -> None: + self.setup_simulation("cpu") + + +class TestRigidConstraintCUDA(BaseRigidConstraintTest): + def setup_method(self) -> None: + self.setup_simulation("cuda") From 032e1ca943ba3f1df8b6894ae5c952ac5d55cc2f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 28 Jun 2026 16:29:40 +0000 Subject: [PATCH 10/16] chore(sim): add __all__ to RigidConstraint module Co-Authored-By: Claude --- embodichain/lab/sim/objects/constraint.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/embodichain/lab/sim/objects/constraint.py b/embodichain/lab/sim/objects/constraint.py index c45fa24d..f52b6468 100644 --- a/embodichain/lab/sim/objects/constraint.py +++ b/embodichain/lab/sim/objects/constraint.py @@ -29,6 +29,8 @@ from embodichain.lab.sim.cfg import RigidConstraintCfg from embodichain.lab.sim.objects.rigid_object import RigidObject +__all__ = ["RigidConstraint"] + @dataclass class RigidConstraint: From b49d81d360308086521030bb57a19887425defb3 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 29 Jun 2026 03:21:35 +0000 Subject: [PATCH 11/16] fix(sim): weld at current relative pose by default With identity local frames the dexsim fixed constraint pulls the two body origins together, which is wrong for the grasp/attach use case. When local_frame_b is None, compute it per env as inv(pose_B) @ pose_A so the constraint welds the objects at their current relative pose. local_frame_a None still defaults to identity. Explicit local_frame_b is used verbatim. Updates the spec, adds two unit tests, and makes the integration test's detach assertion robust (relative-pose drift instead of absolute fall). Co-Authored-By: Claude --- .../2026-06-29-rigid-constraint-design.md | 21 ++++--- embodichain/lab/sim/sim_manager.py | 22 +++++-- tests/sim/objects/test_rigid_constraint.py | 58 ++++++++++++++++++- .../sim/test_rigid_constraint_integration.py | 16 +++-- 4 files changed, 100 insertions(+), 17 deletions(-) diff --git a/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md b/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md index b1d5c07c..3d25ed39 100644 --- a/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md +++ b/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md @@ -63,7 +63,7 @@ FUNCTOR LAYER (gym, on-demand) 1. **One source of truth** — the sim layer owns the constraint registry and all dexsim calls. The functor is a thin adapter: resolve `SceneEntityCfg` → `RigidObject`, then call the sim API. 2. **Per-arena batch symmetry** — `RigidConstraint` mirrors `RigidObject`: N arenas → N dexsim constraint handles, so `env_id i` ↔ arena `i` ↔ handle `i`. Attach/detach can target a subset of `env_ids`. 3. **Fixed-first, extensible** — v1 wires only `create_fixed_constraint`; `RigidConstraintCfg.constraint_type` is reserved (`"fixed"` default). -4. **Local frames default to identity** → attaches at the objects' *current* relative pose. Caller can pass 4×4 or `(N,4,4)`. +4. **Local frames default to the current relative pose** — `local_frame_a` defaults to identity (object A's origin); `local_frame_b` defaults to `inv(pose_B) @ pose_A` (computed from the objects' current poses), so the constraint welds the objects where they are rather than pulling their origins together. Caller can pass explicit 4×4 or `(N,4,4)` matrices to define a specific joint frame. ### Why `SimulationManager`, not `RigidObject`, owns the API @@ -215,7 +215,9 @@ task → event_manager.apply(mode="attach", env_ids=gripping_env_ids) → env.sim.create_rigid_constraint(cfg, env_ids) resolve obj_a, obj_b from self._rigid_objects (raise if missing) target_env_ids = env_ids or range(num_envs) - frames = broadcast(cfg.local_frame_a) # None→I, 4x4→repeat, (N,4,4)→index + frames_a = broadcast(cfg.local_frame_a) # None→I, 4x4→repeat, (N,4,4)→index + frames_b = broadcast(cfg.local_frame_b) if cfg.local_frame_b is not None + else inv(pose_B) @ pose_A per env # default: weld at current relative pose for i in target_env_ids: arena = self.get_env(i) name_i = cfg.name if num_envs==1 else f"{cfg.name}_{i}" @@ -254,16 +256,21 @@ v1 lifecycle is `create` → `remove`. Re-attaching after a partial remove requi At create time the `constraint_handles` list is pre-sized to `num_envs` filled with `None`, then only the `target_env_ids` entries are populated, so arena-index alignment always holds. -### Local-frame broadcasting (once, at create time) +### Local-frame resolution (once, at create time) -| Input | Normalized per-env | +`local_frame_a` is broadcast as below. `local_frame_b` is handled differently +when `None`: instead of identity, it is computed per env as +`inv(pose_B) @ pose_A` from the objects' current poses, so the default welds the +objects at their current relative pose (rather than pulling their origins +together). An explicit `local_frame_b` is broadcast like `local_frame_a`. + +| Input (`local_frame_a`, or explicit `local_frame_b`) | Normalized per-env | |-------|--------------------| -| `None` | `np.eye(4)` for all envs → weld at current relative pose | +| `None` (`local_frame_a`) | `np.eye(4)` for all envs | +| `None` (`local_frame_b`) | `inv(pose_B) @ pose_A` per env (current relative pose) | | `(4, 4)` | same matrix for all envs | | `(N, 4, 4)` | `frames[i]` for env `i` (requires N == num_envs) | -`local_frame_b` mirrors this independently. - ### Interaction with reset Constraints are **not** auto-reset by `reset_objects_state` (constraints aren't bodies). Default task policy: register a `reset`-mode `remove_rigid_constraint` (or call `sim.remove_rigid_constraint(name)` in the task's `_reset`) so stale constraints don't leak across episodes. The sim layer does not silently create/destroy on reset. diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 07c7cbad..be4a7423 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -91,7 +91,7 @@ ) 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", @@ -1115,13 +1115,25 @@ def create_rigid_constraint( else: target_env_ids = list(env_ids) - # broadcast local frames + # 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 ) - frames_b = self._broadcast_frame( - cfg.local_frame_b, 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 diff --git a/tests/sim/objects/test_rigid_constraint.py b/tests/sim/objects/test_rigid_constraint.py index 1fe8e9ca..f866fc8d 100644 --- a/tests/sim/objects/test_rigid_constraint.py +++ b/tests/sim/objects/test_rigid_constraint.py @@ -285,11 +285,21 @@ def get_env(self, arena_index=-1): _broadcast_frame = staticmethod(SimulationManager._broadcast_frame) -def _register_object(sim, uid, num_envs): +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 @@ -498,3 +508,49 @@ def test_partial_remove_then_all_drops_registry(): 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 index 2005f5f5..0b927e5d 100644 --- a/tests/sim/test_rigid_constraint_integration.py +++ b/tests/sim/test_rigid_constraint_integration.py @@ -58,6 +58,12 @@ def _can_run_sim(device: str) -> bool: 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( @@ -122,14 +128,16 @@ def test_fixed_constraint_holds_relative_pose(self): 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 they separate (both fall independently). + # 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 - z_before = float(self.duck_a.get_local_pose(to_matrix=True)[0, 2, 3]) + held_delta_z = self._delta_z() for _ in range(120): self.sim.update(step=1) - z_after = float(self.duck_a.get_local_pose(to_matrix=True)[0, 2, 3]) - assert z_after < z_before - 0.05 + final_delta_z = self._delta_z() + assert abs(final_delta_z - held_delta_z) > 2e-2 class TestRigidConstraintCPU(BaseRigidConstraintTest): From 81de9a014efab40f2918a791ad6b6f064e38141b Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 29 Jun 2026 03:21:35 +0000 Subject: [PATCH 12/16] docs: add rigid-constraint tutorial Two-cube tutorial demonstrating create_rigid_constraint / remove via the SimulationManager API, with a runnable script (CubeCfg, no asset file needed) registered in the tutorial index. Prints the bodies' relative z while attached (held constant) and after removal (free to drift). Co-Authored-By: Claude --- docs/source/tutorial/index.rst | 24 +-- docs/source/tutorial/rigid_constraint.rst | 158 ++++++++++++++++ .../tutorials/sim/create_rigid_constraint.py | 176 ++++++++++++++++++ 3 files changed, 347 insertions(+), 11 deletions(-) create mode 100644 docs/source/tutorial/rigid_constraint.rst create mode 100644 scripts/tutorials/sim/create_rigid_constraint.py 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/scripts/tutorials/sim/create_rigid_constraint.py b/scripts/tutorials/sim/create_rigid_constraint.py new file mode 100644 index 00000000..cec87009 --- /dev/null +++ b/scripts/tutorials/sim/create_rigid_constraint.py @@ -0,0 +1,176 @@ +# ---------------------------------------------------------------------------- +# 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. +""" + +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 --------------- + # local_frame_a / local_frame_b default to identity, so the constraint + # welds the cubes at their *current* relative pose. + 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.") + + 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() From c254c9456ad42b375c27812c8400f762be36e786 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 29 Jun 2026 03:23:01 +0000 Subject: [PATCH 13/16] chore: add from __future__ import annotations to tutorial script Co-Authored-By: Claude --- scripts/tutorials/sim/create_rigid_constraint.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/tutorials/sim/create_rigid_constraint.py b/scripts/tutorials/sim/create_rigid_constraint.py index cec87009..9bee506a 100644 --- a/scripts/tutorials/sim/create_rigid_constraint.py +++ b/scripts/tutorials/sim/create_rigid_constraint.py @@ -19,6 +19,8 @@ observe the constraint holding their relative pose, and then remove it. """ +from __future__ import annotations + import argparse import sys From 4ed0da59cda9dcee64d588d2bc5ccbc595a7a748 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 29 Jun 2026 15:44:29 +0800 Subject: [PATCH 14/16] wip --- .../plans/2026-06-29-rigid-constraint.md | 1691 ----------------- .../2026-06-29-rigid-constraint-design.md | 339 ---- .../tutorials/sim/create_rigid_constraint.py | 4 + 3 files changed, 4 insertions(+), 2030 deletions(-) delete mode 100644 docs/superpowers/plans/2026-06-29-rigid-constraint.md delete mode 100644 docs/superpowers/specs/2026-06-29-rigid-constraint-design.md diff --git a/docs/superpowers/plans/2026-06-29-rigid-constraint.md b/docs/superpowers/plans/2026-06-29-rigid-constraint.md deleted file mode 100644 index d7126bd8..00000000 --- a/docs/superpowers/plans/2026-06-29-rigid-constraint.md +++ /dev/null @@ -1,1691 +0,0 @@ -# Rigid Object Constraint Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Add the ability to attach two `RigidObject`s via a fixed physics constraint and remove it, exposed both as a standalone `SimulationManager` API and as on-demand event functors triggered from a task environment. - -**Architecture:** A sim-layer `RigidConstraint` batch wrapper mirrors `RigidObject`'s per-arena pattern, holding one dexsim `FixedConstraint` handle per arena (with `None` where inactive, so arena-index == list-index). `SimulationManager` owns the constraint registry and all dexsim calls. Two function-style event functors in `events.py` (`create_rigid_constraint`, `remove_rigid_constraint`) resolve `SceneEntityCfg` → `RigidObject` and delegate to the sim API, triggered via custom event modes (`"attach"`/`"detach"`). - -**Tech Stack:** Python 3.11, `embodichain` package, `dexsim` (Warp/PhysX), `torch`, `numpy`, `pytest`, `black==26.3.1`. - -## Global Constraints - -- Every source file begins with the Apache 2.0 copyright header (see `CLAUDE.md`). -- `from __future__ import annotations` at the top of every new file. -- Use `TYPE_CHECKING` guard for `EmbodiedEnv` / `SimulationManager` imports to avoid circular imports. -- Prefer `A | B` over `Union[A, B]`. -- `@configclass` for all config objects; `MISSING` for required fields. -- `logger.log_error(msg)` raises `RuntimeError` by default — error paths call it and let it raise (matches `add_rigid_object`). -- Formatter: `black==26.3.1`. Run `black ` before each commit. -- Package name is `embodichain` (lowercase). -- Spec: `docs/superpowers/specs/2026-06-29-rigid-constraint-design.md`. - ---- - -## File Structure - -``` -embodichain/lab/sim/cfg.py # MODIFY: add RigidConstraintCfg -embodichain/lab/sim/objects/constraint.py # CREATE: RigidConstraint wrapper -embodichain/lab/sim/objects/__init__.py # MODIFY: export RigidConstraint -embodichain/lab/sim/sim_manager.py # MODIFY: +_constraints registry, - # create/remove/get_rigid_constraint, - # asset_uids + _deferred_destroy wiring -embodichain/lab/gym/envs/managers/events.py # MODIFY: +2 functors (no __all__ exists) -tests/sim/objects/test_rigid_constraint.py # CREATE: sim-layer unit tests (mocks) -tests/gym/envs/managers/test_event_rigid_constraint.py # CREATE: functor unit tests (mocks) -tests/sim/test_rigid_constraint_integration.py # CREATE: real-sim smoke (gpu-marked) -``` - -`RigidConstraintCfg` lives in `cfg.py` (where all sim cfgs live). `RigidConstraint` lives in `objects/constraint.py`. The functors live in `events.py` (no `__all__` in that file — just add the functions). - ---- - -### Task 1: `RigidConstraintCfg` in `cfg.py` - -**Files:** -- Modify: `embodichain/lab/sim/cfg.py` (append a new `@configclass` near `RigidObjectGroupCfg` ~line 900) -- Test: `tests/sim/objects/test_rigid_constraint.py` (create file, first test) - -**Interfaces:** -- Produces: `RigidConstraintCfg` dataclass with fields `name: str = MISSING`, `rigid_object_a_uid: str = MISSING`, `rigid_object_b_uid: str = MISSING`, `local_frame_a: np.ndarray | None = None`, `local_frame_b: np.ndarray | None = None`, `constraint_type: Literal["fixed"] = "fixed"`. - -- [ ] **Step 1: Write the failing test** - -Create `tests/sim/objects/test_rigid_constraint.py`: - -```python -# ---------------------------------------------------------------------------- -# 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 - -from dataclasses import MISSING - -from embodichain.lab.sim.cfg import RigidConstraintCfg - - -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) -``` - -- [ ] **Step 2: Run test to verify it fails** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` -Expected: FAIL with `ImportError: cannot import name 'RigidConstraintCfg'` - -- [ ] **Step 3: Write minimal implementation** - -Append to `embodichain/lab/sim/cfg.py` (right after the `RigidObjectGroupCfg` class ends, ~line 1000, before the next `@configclass`). Note `Literal` and `Optional` are already imported at the top of `cfg.py` (lines 22): `from typing import Sequence, Union, Dict, Literal, List, Any, Optional`. - -```python -@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`` attaches at the objects' current relative pose (identity). - 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: As :attr:`local_frame_a`, for object B. 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 (current relative pose).""" - - local_frame_b: np.ndarray | None = None - """Local joint frame on object B. None -> identity (current relative pose).""" - - constraint_type: Literal["fixed"] = "fixed" - """Constraint type. Only ``"fixed"`` is supported in v1.""" -``` - -- [ ] **Step 4: Run test to verify it passes** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` -Expected: PASS (3 tests) - -- [ ] **Step 5: Commit** - -```bash -black embodichain/lab/sim/cfg.py tests/sim/objects/test_rigid_constraint.py -git add embodichain/lab/sim/cfg.py tests/sim/objects/test_rigid_constraint.py -git commit -m "feat(sim): add RigidConstraintCfg" -``` - ---- - -### Task 2: `RigidConstraint` wrapper in `objects/constraint.py` - -**Files:** -- Create: `embodichain/lab/sim/objects/constraint.py` -- Modify: `embodichain/lab/sim/objects/__init__.py` (export `RigidConstraint`) -- Test: `tests/sim/objects/test_rigid_constraint.py` (append) - -**Interfaces:** -- Consumes: `RigidConstraintCfg` (Task 1), `RigidObject` (from `embodichain.lab.sim.objects`). -- Produces: `RigidConstraint` class with: - - `__init__(cfg: RigidConstraintCfg, constraint_handles: list, rigid_object_a: RigidObject, rigid_object_b: RigidObject, device: torch.device)` - - `@property num_envs -> int` - - `get_relative_transform(env_ids=None) -> list[np.ndarray]` - - `get_local_pose(actor_index: int, env_ids=None) -> list[np.ndarray]` - - `get_name(env_id: int) -> str` (returns the base name for single env, `f"{base}_{env_id}"` for multi) - - `is_valid(env_ids=None) -> list[bool]` - - `destroy(self, env_ids=None, arena_resolver=None) -> None` — calls `arena.remove_constraint(name_i)` per env; `arena_resolver(i)` returns the arena for env `i`. When all handles become None, the caller (SimulationManager) drops the wrapper. - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/sim/objects/test_rigid_constraint.py`: - -```python -import torch -from unittest.mock import MagicMock - -from embodichain.lab.sim.objects.constraint import RigidConstraint - - -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) -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` -Expected: FAIL with `ImportError: cannot import name 'RigidConstraint'` - -- [ ] **Step 3: Write minimal implementation** - -Create `embodichain/lab/sim/objects/constraint.py`: - -```python -# ---------------------------------------------------------------------------- -# 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 - - -@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. - rigid_object_b: The second RigidObject. - device: The torch device. - """ - - cfg: RigidConstraintCfg - constraint_handles: list[Any] = field(default_factory=list) - rigid_object_a: RigidObject = None - rigid_object_b: RigidObject = 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 -``` - -Then modify `embodichain/lab/sim/objects/__init__.py` — add the import and keep it exported. Add after the `from .gizmo import Gizmo` line (line 29): - -```python -from .constraint import RigidConstraint -``` - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` -Expected: PASS (all tests, including the new ones) - -- [ ] **Step 5: Commit** - -```bash -black embodichain/lab/sim/objects/constraint.py embodichain/lab/sim/objects/__init__.py tests/sim/objects/test_rigid_constraint.py -git add embodichain/lab/sim/objects/constraint.py embodichain/lab/sim/objects/__init__.py tests/sim/objects/test_rigid_constraint.py -git commit -m "feat(sim): add RigidConstraint wrapper" -``` - ---- - -### Task 3: `SimulationManager` constraint registry + create API - -**Files:** -- Modify: `embodichain/lab/sim/sim_manager.py` (add `_constraints` registry in `__init__`, `create_rigid_constraint`, plus helpers) -- Test: `tests/sim/objects/test_rigid_constraint.py` (append sim-layer tests using a mock sim) - -**Interfaces:** -- Consumes: `RigidConstraintCfg` (Task 1), `RigidConstraint` (Task 2). -- Produces: `SimulationManager.create_rigid_constraint(cfg, env_ids=None) -> RigidConstraint`. The method resolves objects from `self._rigid_objects`, broadcasts local frames, and for each target env calls `self.get_env(i).create_fixed_constraint(name_i, obj_a._entities[i], obj_b._entities[i], frame_a, frame_b)`, stores a `RigidConstraint` in `self._constraints[cfg.name]`. - -**Frame broadcasting rules** (a helper `_broadcast_frame`): `None` -> `np.eye(4)` per env; `(4,4)` -> same matrix per env; `(N,4,4)` -> index `i` (requires `N == num_envs`, else `log_error`). - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/sim/objects/test_rigid_constraint.py`: - -```python -from embodichain.lab.sim.sim_manager import SimulationManager -from embodichain.utils import configclass # noqa: F401 (ensures import works) - - -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 - _broadcast_frame = staticmethod(SimulationManager._broadcast_frame) - - -def _register_object(sim, uid, num_envs): - obj = MagicMock() - obj.uid = uid - obj.num_instances = num_envs - obj._entities = [MagicMock(name=f"{uid}_{i}") for i in range(num_envs)] - 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[i][0] == f"weld_{i}" - assert arena.created[i][1] is sim._rigid_objects["cube"]._entities[i] - assert arena.created[i][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") -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py -v -k "create or broadcast"` -Expected: FAIL with `AttributeError: 'SimulationManager' has no attribute 'create_rigid_constraint'` (and `_broadcast_frame`). - -- [ ] **Step 3: Write minimal implementation** - -In `embodichain/lab/sim/sim_manager.py`: - -3a. Add to the imports block near the top (after the existing `from embodichain.lab.sim.cfg import ...` block, ~line 89). Add `RigidConstraintCfg` to that import: - -```python -from embodichain.lab.sim.cfg import ( - RenderCfg, - PhysicsCfg, - MarkerCfg, - GPUMemoryCfg, - WindowRecordCfg, - LightCfg, - RigidObjectCfg, - SoftObjectCfg, - ClothObjectCfg, - RigidObjectGroupCfg, - ArticulationCfg, - RobotCfg, - RigidConstraintCfg, -) -``` - -3b. Add `RigidConstraint` to the objects import (line ~59-67): - -```python -from embodichain.lab.sim.objects import ( - RigidObject, - RigidObjectGroup, - SoftObject, - ClothObject, - Articulation, - Robot, - Light, - RigidConstraint, -) -``` - -3c. In `__init__`, add the registry next to `self._rigid_objects` (line ~285): - -```python - self._rigid_objects: Dict[str, RigidObject] = dict() - self._constraints: Dict[str, RigidConstraint] = dict() -``` - -3d. Add the methods. Place them right after `get_rigid_object_uid_list` (line ~1005, before `add_rigid_object_group`). - -```python - @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)." - ) - - def create_rigid_constraint( - self, - cfg: RigidConstraintCfg, - env_ids: Sequence[int] | 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 - identity (attach at the objects' current relative pose). - - Args: - cfg: The constraint configuration. - env_ids: Target environment indices. 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 - if env_ids is None: - target_env_ids = list(range(num_envs)) - else: - target_env_ids = list(env_ids) - - # broadcast local frames - frames_a = self._broadcast_frame( - cfg.local_frame_a, num_envs, target_env_ids, cfg.name - ) - 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 -``` - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` -Expected: PASS (all tests so far) - -- [ ] **Step 5: Commit** - -```bash -black embodichain/lab/sim/sim_manager.py tests/sim/objects/test_rigid_constraint.py -git add embodichain/lab/sim/sim_manager.py tests/sim/objects/test_rigid_constraint.py -git commit -m "feat(sim): add SimulationManager.create_rigid_constraint" -``` - ---- - -### Task 4: `remove_rigid_constraint` + `get_rigid_constraint` + registry wiring - -**Files:** -- Modify: `embodichain/lab/sim/sim_manager.py` (add remove/get methods; wire `asset_uids` + `_deferred_destroy`) -- Test: `tests/sim/objects/test_rigid_constraint.py` (append) - -**Interfaces:** -- Consumes: `RigidConstraint.destroy` (Task 2), `self._constraints` (Task 3). -- Produces: - - `remove_rigid_constraint(name, env_ids=None) -> bool` — pops (or partially clears) the constraint; calls `constraint.destroy(env_ids, arena_resolver=self.get_env)`. When all handles are None, drops from registry. Returns True if the constraint existed (or still partially exists after subset remove), False if name unknown. - - `get_rigid_constraint(name) -> RigidConstraint | None` - - `get_rigid_constraint_uid_list() -> list[str]` - - `asset_uids` extended to include constraint names. - - `_deferred_destroy` severs `_constraints`. - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/sim/objects/test_rigid_constraint.py`: - -```python -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") - constraint = 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 -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py -v -k "remove or get_rigid"` -Expected: FAIL with `AttributeError: ... has no attribute 'remove_rigid_constraint'`. - -- [ ] **Step 3: Write minimal implementation** - -In `embodichain/lab/sim/sim_manager.py`, add after `create_rigid_constraint` (from Task 3): - -```python - def remove_rigid_constraint( - self, - name: str, - env_ids: Sequence[int] | 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. 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 - - constraint.destroy(env_ids=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 of constraint names. - """ - return list(self._constraints.keys()) -``` - -Then wire `asset_uids` (line ~421). Add constraints to the returned list. In `asset_uids`: - -```python - @property - def asset_uids(self) -> List[str]: - """Get all assets uid in the simulation. - - The assets include lights, sensors, robots, rigid objects and articulations. - - Returns: - List[str]: list of all assets uid. - """ - uid_list = ["default_plane"] - uid_list.extend(list(self._lights.keys())) - uid_list.extend(list(self._sensors.keys())) - uid_list.extend(list(self._robots.keys())) - uid_list.extend(list(self._rigid_objects.keys())) - uid_list.extend(list(self._rigid_object_groups.keys())) - 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 -``` - -Then wire `_deferred_destroy`. In `_deferred_destroy` (~line 2271, the `_sever_wrapper_refs` calls), add: - -```python - _sever_wrapper_refs("_constraints") -``` - -after the `_sever_wrapper_refs("_rigid_object_groups")` line. Also clear `_constraints` in the explicit `clear()` block — find the line `self._rigid_object_groups` clear region and add: - -```python - self._constraints.clear() -``` - -near the other `.clear()` calls at the end of `_deferred_destroy` (after `self._arenas.clear()`). - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py -v` -Expected: PASS (all sim-layer tests) - -- [ ] **Step 5: Commit** - -```bash -black embodichain/lab/sim/sim_manager.py tests/sim/objects/test_rigid_constraint.py -git add embodichain/lab/sim/sim_manager.py tests/sim/objects/test_rigid_constraint.py -git commit -m "feat(sim): add remove/get_rigid_constraint and registry wiring" -``` - ---- - -### Task 5: Event functors in `events.py` - -**Files:** -- Modify: `embodichain/lab/gym/envs/managers/events.py` (add `create_rigid_constraint` + `remove_rigid_constraint` functions) -- Test: `tests/gym/envs/managers/test_event_rigid_constraint.py` (create) - -**Interfaces:** -- Consumes: `SceneEntityCfg` (from `.cfg`), `RigidConstraintCfg` (from `embodichain.lab.sim.cfg`), `RigidObject` (for isinstance check), `logger`. -- Produces: two module-level functions in `events.py`: - - `create_rigid_constraint(env, env_ids, obj_a_cfg: SceneEntityCfg, obj_b_cfg: SceneEntityCfg, name: str, local_frame_a=None, local_frame_b=None) -> None` - - `remove_rigid_constraint(env, env_ids, name: str) -> None` - -- [ ] **Step 1: Write the failing tests** - -Create `tests/gym/envs/managers/test_event_rigid_constraint.py`: - -```python -# ---------------------------------------------------------------------------- -# 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.lab.gym.envs.managers.cfg import SceneEntityCfg -from embodichain.lab.gym.envs.managers.events import ( - create_rigid_constraint, - remove_rigid_constraint, -) -from embodichain.lab.sim.objects.rigid_object import RigidObject - - -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.""" - env = MagicMock() - env.device = torch.device("cpu") - env.num_envs = 4 - - obj_a = MagicMock() - obj_b = MagicMock() - env.sim.get_asset.side_effect = lambda uid: {"cube": obj_a, "block": obj_b}[uid] - - # control the isinstance check by patching RigidObject for the test - env._obj_a_is_rigid = obj_a_is_rigid - env._obj_b_is_rigid = obj_b_is_rigid - 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).""" - env, obj_a, obj_b = _make_env() - # patch isinstance check to return False for obj_a - monkeypatch.setattr( - "embodichain.lab.gym.envs.managers.events.RigidObject", MagicMock(__instancecheck__=lambda self, o: False) - ) - 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) -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/gym/envs/managers/test_event_rigid_constraint.py -v` -Expected: FAIL with `ImportError: cannot import name 'create_rigid_constraint'` - -- [ ] **Step 3: Write minimal implementation** - -In `embodichain/lab/gym/envs/managers/events.py`, add the imports near the top (the file already imports `RigidObject` at line 28, `logger`, `np`-equivalent math helpers, and `SceneEntityCfg` at line 35). Add to the existing `from embodichain.lab.sim.objects import (...)` block a new import is not needed (RigidObject already there). Add `RigidConstraintCfg` import. Find the `from embodichain.lab.sim.cfg import RigidObjectCfg, ArticulationCfg` line (line 33) and extend it: - -```python -from embodichain.lab.sim.cfg import RigidObjectCfg, ArticulationCfg, RigidConstraintCfg -``` - -Then add `numpy` import if not present (the file uses math helpers but check). At the top, after `import random` (line 21), ensure `import numpy as np` exists. If not present, add it. - -Then append the two functors at the end of the file (after `set_detached_uids_for_env_reset`): - -```python -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 attaches at the - objects' current relative pose. Accepts (4,4) or (N,4,4). - local_frame_b: Local joint frame on object B. None -> identity. - - 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) -``` - -Note: `EmbodiedEnv` is already imported under `TYPE_CHECKING` at the top of `events.py` (line 50). `torch` is already imported (line 19). - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/gym/envs/managers/test_event_rigid_constraint.py -v` -Expected: PASS (5 tests) - -- [ ] **Step 5: Commit** - -```bash -black embodichain/lab/gym/envs/managers/events.py tests/gym/envs/managers/test_event_rigid_constraint.py -git add embodichain/lab/gym/envs/managers/events.py tests/gym/envs/managers/test_event_rigid_constraint.py -git commit -m "feat(env): add rigid-constraint event functors" -``` - ---- - -### Task 6: EventManager custom-mode wiring test - -**Files:** -- Test: `tests/gym/envs/managers/test_event_rigid_constraint.py` (append) - -**Interfaces:** -- Consumes: `EventManager`, `EventCfg`, the two functors (Task 5), `ManagerBase`. -- Produces: a test proving `EventManager.apply(mode="attach", env_ids)` invokes a registered custom-mode functor once with those env_ids (no env subclassing needed — the manager supports arbitrary mode strings). - -- [ ] **Step 1: Write the failing test** - -Append to `tests/gym/envs/managers/test_event_rigid_constraint.py`: - -```python -from embodichain.lab.gym.envs.managers.event_manager import EventManager -from embodichain.lab.gym.envs.managers.cfg import EventCfg -from embodichain.utils import configclass - - -@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() - 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 -``` - -- [ ] **Step 2: Run test to verify it fails** - -Run: `pytest tests/gym/envs/managers/test_event_rigid_constraint.py::test_custom_mode_apply_invokes_functor_with_env_ids -v` -Expected: FAIL — likely an import or attribute error if `EventManager` construction with the mock env fails. If the mock env is insufficient, fix the mock (the test asserts behavior, so adjust the mock to satisfy EventManager's needs). Inspect the failure and patch the mock accordingly (e.g. add `env.cfg = None` if needed). - -- [ ] **Step 3: (No new implementation needed)** - -`EventManager` already supports arbitrary mode strings (see `event_manager.py` `_prepare_functors` — it registers any `functor_cfg.mode` into `_mode_functor_names`). If the test fails on a missing mock attribute, fix the test's mock rather than changing `EventManager`. Document what was needed. - -- [ ] **Step 4: Run test to verify it passes** - -Run: `pytest tests/gym/envs/managers/test_event_rigid_constraint.py -v` -Expected: PASS (6 tests) - -- [ ] **Step 5: Commit** - -```bash -black tests/gym/envs/managers/test_event_rigid_constraint.py -git add tests/gym/envs/managers/test_event_rigid_constraint.py -git commit -m "test(env): custom-mode apply invokes rigid-constraint functor" -``` - ---- - -### Task 7: Real-sim integration smoke test - -**Files:** -- Create: `tests/sim/test_rigid_constraint_integration.py` - -**Interfaces:** -- Consumes: `SimulationManager`, `RigidObjectCfg`, `RigidConstraintCfg`, real dexsim. Mirrors the contract in `dexsim/python/test/engine/test_constraint.py` at the EmbodiChain layer. -- Produces: a skipped-unless-gpu test that attaches two dynamic cubes, steps, asserts the relative transform stays constant, detaches, steps, asserts separation. - -- [ ] **Step 1: Write the test** - -Create `tests/sim/test_rigid_constraint_integration.py`: - -```python -# ---------------------------------------------------------------------------- -# 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. - -Skipped unless a GPU/display is available. Mirrors the dexsim -test_constraint.py contract at the EmbodiChain SimulationManager layer. -""" - -from __future__ import annotations - -import numpy as np -import pytest -import torch - -from embodichain.lab.sim.sim_manager import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import ( - RigidObjectCfg, - RigidConstraintCfg, - RigidBodyAttributesCfg, -) -from embodichain.lab.sim.shapes import MeshCfg -from dexsim.types import ActorType, RigidBodyShape - - -def _can_run_gpu_sim() -> bool: - try: - return torch.cuda.is_available() - except Exception: - return False - - -pytestmark = pytest.mark.skipif( - not _can_run_gpu_sim(), reason="GPU simulation required for constraint integration test" -) - - -def _make_sim(num_envs: int = 1) -> SimulationManager: - cfg = SimulationManagerCfg() - cfg.headless = True - cfg.sim_device = "cuda" - cfg.num_envs = num_envs - return SimulationManager(cfg) - - -def test_fixed_constraint_holds_relative_pose(): - """Two welded cubes keep their relative transform under physics; detach lets them separate.""" - sim = _make_sim(num_envs=1) - try: - attrs = RigidBodyAttributesCfg() - attrs.mass = 0.2 - cube_cfg = RigidObjectCfg( - uid="cube_a", - body_type="dynamic", - init_pos=[0.0, 0.0, 1.4], - attrs=attrs, - shape=MeshCfg(fpath="..."), # placeholder; see note - ) - block_cfg = RigidObjectCfg( - uid="cube_b", - body_type="dynamic", - init_pos=[0.0, 0.0, 1.2], - attrs=attrs, - shape=MeshCfg(fpath="..."), - ) - cube = sim.add_rigid_object(cube_cfg) - block = sim.add_rigid_object(block_cfg) - - constraint = sim.create_rigid_constraint( - RigidConstraintCfg( - name="weld", - rigid_object_a_uid="cube_a", - rigid_object_b_uid="cube_b", - ) - ) - assert constraint.is_valid() == [True] - - initial_delta_z = ( - block.get_local_pose()[0, 2, 3] - cube.get_local_pose()[0, 2, 3] - ) - - for _ in range(120): - sim.update(step=1) - - rel = constraint.get_relative_transform()[0] - np.testing.assert_allclose(rel[:3, 3], np.zeros(3), atol=2e-2) - delta_z = ( - block.get_local_pose()[0, 2, 3] - cube.get_local_pose()[0, 2, 3] - ) - assert abs(delta_z - initial_delta_z) < 2e-2 - - # detach and confirm they separate - sim.remove_rigid_constraint("weld") - z_before = cube.get_local_pose()[0, 2, 3] - for _ in range(120): - sim.update(step=1) - z_after = cube.get_local_pose()[0, 2, 3] - assert z_after < z_before - 0.05 - finally: - sim.destroy(exit_process=False) - SimulationManager.flush_cleanup_queue() -``` - -- [ ] **Step 2: Run test to verify it skips (or passes on GPU)** - -Run: `pytest tests/sim/test_rigid_constraint_integration.py -v` -Expected: SKIPPED (no GPU) on CPU CI. On a GPU machine it should PASS. - -- [ ] **Step 3: (Implementation already complete from Tasks 1-5)** - -This test exercises the full sim path. If `RigidObjectCfg`/`MeshCfg` construction needs a real mesh path, replace the `fpath="..."` placeholders with a real asset path from `embodichain/data/assets` (check `SimResources` for a built-in cube mesh). Document the chosen path in the test. If the test cannot run headless on the CI runner, leave it gpu-marked and skipped — that's acceptable per the spec. - -- [ ] **Step 4: Commit** - -```bash -black tests/sim/test_rigid_constraint_integration.py -git add tests/sim/test_rigid_constraint_integration.py -git commit -m "test(sim): add rigid-constraint real-sim integration smoke test" -``` - ---- - -### Task 8: Run full suite, black, finalize - -**Files:** -- All touched files. - -- [ ] **Step 1: Run the full constraint test suite** - -Run: `pytest tests/sim/objects/test_rigid_constraint.py tests/gym/envs/managers/test_event_rigid_constraint.py tests/sim/test_rigid_constraint_integration.py -v` -Expected: all unit tests PASS; integration test SKIPPED (or PASS on GPU). - -- [ ] **Step 2: Run black on all changed files** - -Run: `black embodichain/lab/sim/cfg.py embodichain/lab/sim/objects/constraint.py embodichain/lab/sim/objects/__init__.py embodichain/lab/sim/sim_manager.py embodichain/lab/gym/envs/managers/events.py tests/sim/objects/test_rigid_constraint.py tests/gym/envs/managers/test_event_rigid_constraint.py tests/sim/test_rigid_constraint_integration.py` -Expected: no changes (or only formatting). - -- [ ] **Step 3: Run the pre-commit check skill** - -Use the `/pre-commit-check` skill (it runs all local CI checks). Fix any violations it reports. - -- [ ] **Step 4: Commit any formatting fixes** - -```bash -git add -A -git commit -m "chore: black formatting + pre-commit fixes" -``` - -- [ ] **Step 5: Final summary** - -The implementation is complete. Summary of what was built: -- `RigidConstraintCfg` (cfg.py) -- `RigidConstraint` wrapper (objects/constraint.py) -- `SimulationManager.create_rigid_constraint` / `remove_rigid_constraint` / `get_rigid_constraint` + `_constraints` registry -- `create_rigid_constraint` / `remove_rigid_constraint` event functors (events.py) -- Unit tests (mocks) + integration smoke test (gpu-marked) - -## Self-Review - -**1. Spec coverage:** -- §3 Architecture (sim layer + functor layer): Tasks 1-5. -- §4 Sim-layer API (RigidConstraint, RigidConstraintCfg, SimulationManager methods): Tasks 1-4. -- §5 Functor layer (two functors, registration/triggering): Tasks 5-6. -- §6 Data flow (create/remove, per-env selectivity, frame broadcasting, reset interaction): Tasks 3-4 + Task 6 (reset interaction is task-policy, documented in spec; no sim code needed). -- §7 Error handling: each error case is a test in Task 3-4. -- §8 Testing: Tasks 3 (sim unit), 5-6 (functor unit), 7 (integration). -- §9 File layout: matches. - -**2. Placeholder scan:** The `MeshCfg(fpath="...")` in Task 7 is flagged as needing a real asset path — that's an integration-test asset-resolution detail, addressed in Task 7 Step 3. No "TBD"/"implement later" in deliverable code. - -**3. Type consistency:** `create_rigid_constraint` / `remove_rigid_constraint` / `get_rigid_constraint` names match across sim layer (Task 3-4) and functor layer (Task 5). `RigidConstraint` field names (`constraint_handles`, `rigid_object_a`, `rigid_object_b`, `device`) match Task 2 init and Task 3 usage. `destroy(env_ids, arena_resolver)` signature matches Task 2 and Task 4's call. `_broadcast_frame` is a `@staticmethod` accessed as `SimulationManager._broadcast_frame` in tests (Task 3) and as `self._broadcast_frame` in `create_rigid_constraint` — both valid. diff --git a/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md b/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md deleted file mode 100644 index 3d25ed39..00000000 --- a/docs/superpowers/specs/2026-06-29-rigid-constraint-design.md +++ /dev/null @@ -1,339 +0,0 @@ -# Rigid Object Constraint — Design Spec - -- **Date:** 2026-06-29 -- **Status:** Approved -- **Topic:** Fixed constraint support between two `RigidObject`s, plus an on-demand event functor - -## 1. Goal - -Add the ability to attach two rigid objects via a fixed physics constraint and to remove that constraint again, both as a standalone simulation API (usable outside the gym) and as an on-demand event functor triggered from a task environment. - -### Functional requirements - -1. Two `RigidObject`s can be attached via a constraint. -2. The constraint between the two rigid objects can be removed. -3. A functor creates and removes the constraint inside a task environment. - -### Non-goals (v1, deferred with extension points) - -- Prismatic / revolute / spherical / d6 typed constraints (reserved `constraint_type` field). -- Per-constraint physics tuning (limits, drive, motion) — comes with typed constraints. -- Articulation ↔ rigid and articulation ↔ articulation constraints — `RigidObject` ↔ `RigidObject` only. -- Auto-detach-on-reset baked into the sim layer — kept as task policy. - -## 2. Reference - -The design mirrors the dexsim `fixed_constraint` example (`dexsim/examples/python/physics/rigidbody/fixed_constraint.py`) and lifts its API onto EmbodiChain's batched object/manager layer. - -dexsim constraint API (bound on `Arena`, inherited by `Env`; C++ in `dexsim/cpp/pybind/environment/environment.cpp`): - -- `arena.create_fixed_constraint(name, actor0, actor1, local_frame0=I, local_frame1=I) -> FixedConstraint | None` -- `arena.remove_constraint(name)` -- `arena.get_constraint(name) -> Constraint | None` -- `arena.get_all_constraints() -> list[Constraint]` -- Constraint handle exposes `get_name()`, `get_constraint_type()`, `is_valid()`, `get_local_pose(idx)`, `get_relative_transform()`. -- Typed variants exist (`create_prismatic_constraint`, `create_revolute_constraint`, `create_spherical_constraint`, `create_d6_constraint`) — reserved for later. - -## 3. Architecture - -Two layers: - -``` -SIM LAYER (standalone, usable without gym) - SimulationManager RigidConstraint (new) - ├─ create_rigid_constraint(...) ──► batch wrapper over N arenas - ├─ remove_rigid_constraint(name) ├─ per-arena dexsim handles - ├─ get_rigid_constraint(name) ├─ obj_a / obj_b refs - └─ self._constraints: dict ├─ get_relative_transform() - {name: RigidConstraint} ├─ get_local_pose(idx) - └─ destroy() - delegates to dexsim Arena.create_fixed_constraint / - Arena.remove_constraint (per env_id / arena) - -FUNCTOR LAYER (gym, on-demand) - events.py - ├─ create_rigid_constraint(env, env_ids, obj_a_cfg, obj_b_cfg, name, ...) - └─ remove_rigid_constraint(env, env_ids, name) - registered under EventCfg(mode=, ...); - task triggers via env.event_manager.apply(mode="attach", env_ids) -``` - -### Principles - -1. **One source of truth** — the sim layer owns the constraint registry and all dexsim calls. The functor is a thin adapter: resolve `SceneEntityCfg` → `RigidObject`, then call the sim API. -2. **Per-arena batch symmetry** — `RigidConstraint` mirrors `RigidObject`: N arenas → N dexsim constraint handles, so `env_id i` ↔ arena `i` ↔ handle `i`. Attach/detach can target a subset of `env_ids`. -3. **Fixed-first, extensible** — v1 wires only `create_fixed_constraint`; `RigidConstraintCfg.constraint_type` is reserved (`"fixed"` default). -4. **Local frames default to the current relative pose** — `local_frame_a` defaults to identity (object A's origin); `local_frame_b` defaults to `inv(pose_B) @ pose_A` (computed from the objects' current poses), so the constraint welds the objects where they are rather than pulling their origins together. Caller can pass explicit 4×4 or `(N,4,4)` matrices to define a specific joint frame. - -### Why `SimulationManager`, not `RigidObject`, owns the API - -The sim owns `self._arenas`, `self._env`, and every scene-mutation method (`add_rigid_object`, `remove_asset`, `draw_marker`). A constraint lives *between* two objects, so it belongs to the scene owner, not to either object. This keeps `RigidObject` focused on a single body and the constraint registry co-located with the rigid-object registry. - -## 4. Sim-layer API - -### New file: `embodichain/lab/sim/objects/constraint.py` - -```python -@dataclass -class RigidConstraint: - """Batch of fixed constraints linking two RigidObjects across all arenas. - - Each entry binds rigid_object_a's entity[i] to rigid_object_b's entity[i] - within arena[i] via a dexsim FixedConstraint. - """ - cfg: RigidConstraintCfg - constraint_handles: list[Any] # length == num_envs; None where inactive - rigid_object_a: RigidObject - rigid_object_b: RigidObject - device: torch.device - - @property - def num_envs(self) -> int: ... - - def get_relative_transform(self, env_ids=None) -> list[np.ndarray]: ... - def get_local_pose(self, actor_index: int, env_ids=None) -> list[np.ndarray]: ... - def get_name(self, env_id: int) -> str: ... - def is_valid(self, env_ids=None) -> list[bool]: ... - def destroy(self, env_ids: Sequence[int] | None = None) -> None: ... -``` - -`constraint_handles` is a list of length `num_envs` with `None` wherever the constraint is not active in that arena, so **arena index == list index** always holds. - -### `RigidConstraintCfg` (in `embodichain/lab/sim/cfg.py`) - -```python -@configclass -class RigidConstraintCfg: - name: str = MISSING - rigid_object_a_uid: str = MISSING - rigid_object_b_uid: str = MISSING - local_frame_a: np.ndarray | None = None # None → identity; 4x4 or (N,4,4) - local_frame_b: np.ndarray | None = None - constraint_type: Literal["fixed"] = "fixed" # reserved for typed constraints -``` - -### `SimulationManager` additions - -```python -def create_rigid_constraint( - self, cfg: RigidConstraintCfg, env_ids: Sequence[int] | None = None -) -> RigidConstraint: - """Create a fixed constraint between two RigidObjects (env_ids-aware).""" - -def remove_rigid_constraint( - self, name: str, env_ids: Sequence[int] | None = None -) -> bool: - """Remove by base name; idempotent. env_ids subset clears only those arenas.""" - -def get_rigid_constraint(self, name: str) -> RigidConstraint | None: ... -def get_rigid_constraint_uid_list(self) -> list[str]: ... -``` - -- New registry `self._constraints: Dict[str, RigidConstraint]`. -- `_deferred_destroy` severs `_constraints` like the other registries. - -## 5. Functor layer - -Two function-style event functors in `embodichain/lab/gym/envs/managers/events.py`, following the `add-functor` conventions: signature `(env, env_ids, ...) -> None`, `SceneEntityCfg` for entity refs, `from __future__ import annotations`, `TYPE_CHECKING` guard for `EmbodiedEnv`. - -### `create_rigid_constraint` (attach) - -```python -def create_rigid_constraint( - env, env_ids, - 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.""" - obj_a = env.sim.get_asset(obj_a_cfg.uid) - obj_b = env.sim.get_asset(obj_b_cfg.uid) - # type-check both are RigidObject; else log_error - 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, - ) -``` - -### `remove_rigid_constraint` (detach) - -```python -def remove_rigid_constraint(env, env_ids, name: str) -> None: - """Remove the named constraint for the given env_ids. Idempotent.""" - env.sim.remove_rigid_constraint(name, env_ids=env_ids) -``` - -### Registration & triggering - -```python -@configclass -class MyTaskEventsCfg: - attach_objects: EventCfg = EventCfg( - func=create_rigid_constraint, mode="attach", - params={ - "obj_a_cfg": SceneEntityCfg(uid="cube"), - "obj_b_cfg": SceneEntityCfg(uid="block"), - "name": "cube_block_weld", - }, - ) - detach_objects: EventCfg = EventCfg( - func=remove_rigid_constraint, mode="detach", - params={"name": "cube_block_weld"}, - ) -``` - -```python -self.event_manager.apply(mode="attach", env_ids=gripping_env_ids) -self.event_manager.apply(mode="detach", env_ids=released_env_ids) -``` - -### Decisions - -1. **Thin adapter** — functor does resolution + delegation only; no dexsim calls, no state. -2. **`env_ids` threading** — forwarded to `env.sim.create_rigid_constraint(..., env_ids=...)`. -3. **Custom modes** (`"attach"`/`"detach"`) — task-driven; supported by `EventManager` (arbitrary mode string, task wires `apply`). -4. **Detach takes only `name` + `env_ids`** — looked up by name in the sim registry. - -## 6. Data flow - -### Create (attach) - -``` -task → event_manager.apply(mode="attach", env_ids=gripping_env_ids) - → EventManager iterates "attach" functors → func(env, env_ids, **params) - → create_rigid_constraint(env, env_ids, obj_a_cfg, obj_b_cfg, name, frames) - obj_a = env.sim.get_asset(obj_a_cfg.uid); obj_b = env.sim.get_asset(obj_b_cfg.uid) - build RigidConstraintCfg(...) - → env.sim.create_rigid_constraint(cfg, env_ids) - resolve obj_a, obj_b from self._rigid_objects (raise if missing) - target_env_ids = env_ids or range(num_envs) - frames_a = broadcast(cfg.local_frame_a) # None→I, 4x4→repeat, (N,4,4)→index - frames_b = broadcast(cfg.local_frame_b) if cfg.local_frame_b is not None - else inv(pose_B) @ pose_A per env # default: weld at current relative pose - for i in target_env_ids: - arena = self.get_env(i) - name_i = cfg.name if num_envs==1 else f"{cfg.name}_{i}" - handle = arena.create_fixed_constraint(name_i, obj_a[i], obj_b[i], fa, fb) - if handle is None: log_error(arena i) - self._constraints[cfg.name] = RigidConstraint(...) - → physics: obj_a[i] and obj_b[i] welded in arena[i] -``` - -### Remove (detach) - -``` -task → event_manager.apply(mode="detach", env_ids=released_env_ids) - → remove_rigid_constraint(env, env_ids, name) - → env.sim.remove_rigid_constraint(name, env_ids) - constraint = self._constraints.pop(name, None) - if None: log_warning; return False - constraint.destroy(env_ids) # arena.remove_constraint(name_i) per env - → physics: obj_a[i] and obj_b[i] no longer welded -``` - -### Per-env selectivity — index alignment - -`constraint_handles` is a list of length `num_envs`, `None` wherever inactive: - -``` -num_envs = 4, attach on env_ids=[0, 2] -constraint_handles = [handle_0, None, handle_2, None] -``` - -- `get_relative_transform(env_ids=[2])` reads `constraint_handles[2]`. -- `remove(env_ids=[0])` clears index 0 only; index 2 stays attached. -- Partial remove: `destroy(env_ids=[0])` → `arena.remove_constraint("name_0")` + `constraint_handles[0] = None`. When all handles become `None`, the wrapper is dropped from `self._constraints` (base name freed). - -v1 lifecycle is `create` → `remove`. Re-attaching after a partial remove requires removing the whole constraint by name and creating it again (consistent with `add_rigid_object`'s "already exists" semantics — a duplicate base name on `create` is an error, even if only some envs are currently active). - -At create time the `constraint_handles` list is pre-sized to `num_envs` filled with `None`, then only the `target_env_ids` entries are populated, so arena-index alignment always holds. - -### Local-frame resolution (once, at create time) - -`local_frame_a` is broadcast as below. `local_frame_b` is handled differently -when `None`: instead of identity, it is computed per env as -`inv(pose_B) @ pose_A` from the objects' current poses, so the default welds the -objects at their current relative pose (rather than pulling their origins -together). An explicit `local_frame_b` is broadcast like `local_frame_a`. - -| Input (`local_frame_a`, or explicit `local_frame_b`) | Normalized per-env | -|-------|--------------------| -| `None` (`local_frame_a`) | `np.eye(4)` for all envs | -| `None` (`local_frame_b`) | `inv(pose_B) @ pose_A` per env (current relative pose) | -| `(4, 4)` | same matrix for all envs | -| `(N, 4, 4)` | `frames[i]` for env `i` (requires N == num_envs) | - -### Interaction with reset - -Constraints are **not** auto-reset by `reset_objects_state` (constraints aren't bodies). Default task policy: register a `reset`-mode `remove_rigid_constraint` (or call `sim.remove_rigid_constraint(name)` in the task's `_reset`) so stale constraints don't leak across episodes. The sim layer does not silently create/destroy on reset. - -## 7. Error handling - -| Condition | Layer | Behavior | -|-----------|-------|----------| -| Either RigidObject uid missing | sim | `log_error` (raises) | -| Duplicate base name in `_constraints` | sim | `log_error` | -| `create_fixed_constraint` returns `None` | sim | `log_error` with arena index | -| `(N,4,4)` frame N ≠ num_envs | sim | `log_error` | -| `remove` on unknown name | sim | `log_warning`, return `False` | -| `remove` with env_ids subset | sim | clear only those handles | -| Entity in functor not a `RigidObject` | functor | `log_error` | -| Remove a constraint already removed (handle `None`) | sim | no-op, success | -| Static + dynamic body combo | sim | allowed (weld-to-environment) | - -## 8. Testing - -**`tests/sim/objects/test_rigid_constraint.py`** (sim layer, mocks — `MockSim` exposing `create_fixed_constraint`/`remove_constraint` returning mock handles; `add_rigid_object` returning mock `RigidObject`s with `_entities`): - -- `test_create_resolves_both_objects` -- `test_create_missing_object_raises` -- `test_create_subset_env_ids` — handles at 0,2 only; `None` elsewhere; alignment holds -- `test_local_frame_broadcasting` — None→identity, 4×4→repeat, (N,4,4)→indexed, bad N→error -- `test_remove_by_name_clears_handles` -- `test_remove_unknown_name_warns_false` -- `test_partial_remove_keeps_others` -- `test_all_removed_drops_from_registry` -- `test_get_relative_transform_skips_none_handles` - -**`tests/gym/envs/managers/test_event_rigid_constraint.py`** (functor layer, mocks — `MockEnv` with `MockSim`, spied `create/remove_rigid_constraint`): - -- `test_create_functor_delegates_to_sim` — params forwarded, `env_ids` threaded -- `test_create_functor_rejects_non_rigid_object` — `Articulation` uid → error -- `test_remove_functor_delegates` — forwards name + env_ids -- `test_custom_mode_apply_invokes_functor` — `EventCfg(mode="attach")` → `apply("attach", env_ids)` calls the spy once with those env_ids - -**`tests/sim/test_rigid_constraint_integration.py`** (real-sim smoke, `@pytest.mark.gpu` / skipped without display): - -- Mirror the dexsim `test_constraint.py` contract at the EmbodiChain layer: attach two dynamic cubes, step, assert relative transform stays constant; detach, step, assert they separate. - -## 9. File layout - -``` -embodichain/lab/sim/objects/constraint.py # RigidConstraint -embodichain/lab/sim/objects/__init__.py # export RigidConstraint -embodichain/lab/sim/cfg.py # + RigidConstraintCfg -embodichain/lab/sim/sim_manager.py # + create/remove/get_rigid_constraint, - # _constraints registry, asset_uids + - # _deferred_destroy wiring -embodichain/lab/gym/envs/managers/events.py # + create/remove_rigid_constraint functors - # + __all__ -tests/sim/objects/test_rigid_constraint.py # sim-layer unit (mocks) -tests/gym/envs/managers/test_event_rigid_constraint.py # functor unit (mocks) -tests/sim/test_rigid_constraint_integration.py # real-sim smoke (gpu-marked) -``` - -## 10. Out of scope / extension points - -- Typed constraints (prismatic/revolute/spherical/d6) → `RigidConstraintCfg.constraint_type` reserved; typed factories land later without public API change. -- Per-constraint physics tuning (limits, drive, motion) → with typed constraints. -- Articulation constraints → `RigidObject`↔`RigidObject` only for v1. -- Auto-detach-on-reset → task policy, not sim policy. -- A dedicated `/add-...` skill → not needed; functors follow existing `add-functor` conventions. diff --git a/scripts/tutorials/sim/create_rigid_constraint.py b/scripts/tutorials/sim/create_rigid_constraint.py index 9bee506a..1615265e 100644 --- a/scripts/tutorials/sim/create_rigid_constraint.py +++ b/scripts/tutorials/sim/create_rigid_constraint.py @@ -130,6 +130,10 @@ def main(): 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) From 2fa44eb1f0f46340f7ae9016b146f47314f8fc19 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 29 Jun 2026 07:50:14 +0000 Subject: [PATCH 15/16] Address Copilot review: env_ids type contract + doc/type fixes - Normalize env_ids (None / tensor / sequence) to list[int] at the sim layer via _normalize_env_ids, used in create_rigid_constraint and remove_rigid_constraint. Widens the sim API to accept torch.Tensor (as passed by EventManager) so the per-arena dexsim names are clean ('weld_0' not a tensor stringification) and create/remove agree. Adds tensor-env_ids regression tests. - Correct local_frame_b docstrings/comments across RigidConstraintCfg, the create functor, the tutorial script, and the integration test: None -> inv(pose_B) @ pose_A (weld at current relative pose), not identity. local_frame_a None stays identity. - Make RigidConstraint.rigid_object_a/b optional in the type hint (RigidObject | None) to match their None defaults. - Move mid-file test imports to module top (E402); drop an unused RigidObject import. - Add teardown_method to the integration base test (destroy + flush) to avoid leaking dexsim scenes between tests. Co-Authored-By: Claude --- embodichain/lab/gym/envs/managers/events.py | 8 +-- embodichain/lab/sim/cfg.py | 18 ++++--- embodichain/lab/sim/objects/constraint.py | 8 +-- embodichain/lab/sim/sim_manager.py | 52 +++++++++++++++---- .../tutorials/sim/create_rigid_constraint.py | 6 ++- .../managers/test_event_rigid_constraint.py | 10 ++-- tests/sim/objects/test_rigid_constraint.py | 39 ++++++++++++-- .../sim/test_rigid_constraint_integration.py | 19 ++++++- 8 files changed, 123 insertions(+), 37 deletions(-) diff --git a/embodichain/lab/gym/envs/managers/events.py b/embodichain/lab/gym/envs/managers/events.py index 4122ddca..a531ac40 100644 --- a/embodichain/lab/gym/envs/managers/events.py +++ b/embodichain/lab/gym/envs/managers/events.py @@ -644,9 +644,11 @@ def create_rigid_constraint( 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 attaches at the - objects' current relative pose. Accepts (4,4) or (N,4,4). - local_frame_b: Local joint frame on object B. None -> identity. + 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. diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index ce137e16..4ad97cee 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -1010,10 +1010,15 @@ class RigidConstraintCfg: 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`` attaches at the objects' current relative pose (identity). - 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: As :attr:`local_frame_a`, for object B. Defaults to None. + ``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. @@ -1032,10 +1037,11 @@ class RigidConstraintCfg: """UID of the second RigidObject.""" local_frame_a: np.ndarray | None = None - """Local joint frame on object A. None -> identity (current relative pose).""" + """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 -> identity (current relative pose).""" + """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.""" diff --git a/embodichain/lab/sim/objects/constraint.py b/embodichain/lab/sim/objects/constraint.py index f52b6468..0d8951bb 100644 --- a/embodichain/lab/sim/objects/constraint.py +++ b/embodichain/lab/sim/objects/constraint.py @@ -45,15 +45,15 @@ class RigidConstraint: Args: cfg: The constraint configuration. constraint_handles: Per-arena dexsim constraint handles (None where inactive). - rigid_object_a: The first RigidObject. - rigid_object_b: The second RigidObject. + 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 - rigid_object_b: RigidObject = None + rigid_object_a: RigidObject | None = None + rigid_object_b: RigidObject | None = None device: torch.device = field(default_factory=torch.device("cpu")) @property diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index be4a7423..46340601 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -1047,20 +1047,51 @@ def _broadcast_frame( "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] | None = None, + 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 - identity (attach at the objects' current relative pose). + 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. None -> all arenas. + 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`. @@ -1109,11 +1140,8 @@ def create_rigid_constraint( f"{rigid_object_b.num_instances} instances but num_envs is {num_envs}." ) - # resolve target env_ids - if env_ids is None: - target_env_ids = list(range(num_envs)) - else: - target_env_ids = list(env_ids) + # 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). @@ -1182,7 +1210,7 @@ def get_cloth_object_uid_list(self) -> List[str]: def remove_rigid_constraint( self, name: str, - env_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, ) -> bool: """Remove a rigid constraint by name. @@ -1192,7 +1220,8 @@ def remove_rigid_constraint( Args: name: The base constraint name. - env_ids: Subset of arenas to clear. None -> all. + 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), @@ -1203,7 +1232,8 @@ def remove_rigid_constraint( logger.log_warning(f"Constraint '{name}' not found. Nothing to remove.") return False - constraint.destroy(env_ids=env_ids, arena_resolver=self.get_env) + 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): diff --git a/scripts/tutorials/sim/create_rigid_constraint.py b/scripts/tutorials/sim/create_rigid_constraint.py index 9bee506a..68a140f0 100644 --- a/scripts/tutorials/sim/create_rigid_constraint.py +++ b/scripts/tutorials/sim/create_rigid_constraint.py @@ -107,8 +107,10 @@ def main(): print("[INFO]: Scene setup complete with two cubes (cube_a, cube_b).") # --- Phase 1: attach the two cubes with a fixed constraint --------------- - # local_frame_a / local_frame_b default to identity, so the constraint - # welds the cubes at their *current* relative pose. + # 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", diff --git a/tests/gym/envs/managers/test_event_rigid_constraint.py b/tests/gym/envs/managers/test_event_rigid_constraint.py index b80ac704..fb40a03e 100644 --- a/tests/gym/envs/managers/test_event_rigid_constraint.py +++ b/tests/gym/envs/managers/test_event_rigid_constraint.py @@ -23,12 +23,13 @@ import torch from unittest.mock import MagicMock -from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg +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, ) -from embodichain.lab.sim.objects.rigid_object import RigidObject class MockRigidObjectForFunctor: @@ -143,11 +144,6 @@ def test_remove_functor_none_env_ids(): env.sim.remove_rigid_constraint.assert_called_once_with("weld", env_ids=None) -from embodichain.lab.gym.envs.managers.event_manager import EventManager -from embodichain.lab.gym.envs.managers.cfg import EventCfg -from embodichain.utils import configclass - - @configclass class _AttachEventsCfg: attach: EventCfg = EventCfg( diff --git a/tests/sim/objects/test_rigid_constraint.py b/tests/sim/objects/test_rigid_constraint.py index f866fc8d..55bd9d43 100644 --- a/tests/sim/objects/test_rigid_constraint.py +++ b/tests/sim/objects/test_rigid_constraint.py @@ -27,6 +27,7 @@ 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(): @@ -224,9 +225,6 @@ def test_rigid_constraint_destroy_all_returns_all_cleared(): assert all(h is None for h in constraint.constraint_handles) -from embodichain.lab.sim.sim_manager import SimulationManager - - class MockArena: """Mock dexsim arena that records created constraints.""" @@ -283,6 +281,7 @@ def get_env(self, arena_index=-1): 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): @@ -433,6 +432,40 @@ def test_broadcast_frame_N4x4_indexes(): 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) diff --git a/tests/sim/test_rigid_constraint_integration.py b/tests/sim/test_rigid_constraint_integration.py index 0b927e5d..c0fd27f1 100644 --- a/tests/sim/test_rigid_constraint_integration.py +++ b/tests/sim/test_rigid_constraint_integration.py @@ -74,7 +74,8 @@ def setup_simulation(self, sim_device: str) -> None: self.sim.enable_physics(False) duck_path = get_data_path(DUCK_PATH) - # Two dynamic ducks at different heights, welded at identity frames. + # 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() @@ -102,6 +103,22 @@ def setup_simulation(self, sim_device: str) -> None: 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( From 9339f09cdc85340f60972d1b87dd24e68b3a6524 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 29 Jun 2026 14:21:17 +0000 Subject: [PATCH 16/16] wip --- tests/sim/test_rigid_constraint_integration.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/sim/test_rigid_constraint_integration.py b/tests/sim/test_rigid_constraint_integration.py index c0fd27f1..4cf49c51 100644 --- a/tests/sim/test_rigid_constraint_integration.py +++ b/tests/sim/test_rigid_constraint_integration.py @@ -162,6 +162,7 @@ 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")