feat: consume physical entities in manager reset transactions (#1599)

docs/sphinx/source/changelog.md

@@ -17,6 +17,11 @@ PyPI 版本变更与未发布变更；发布日期采用 PyPI 上传日期。完
 
 ### Added / 新增
 
+- Add the M2 entity consumer: typed physical sources/variants, explicit logical-to-physical binding, selected entity reset transactions and per-variant state/control defaults. A registered primitive task validates the same pickleable EnvFactory with MuJoCo and IsaacSim. Standard and ROCm dependency profiles require the released `unisim-core>=1.5.0` package, without a Git source override. Regression tests use the unified worker and construct the nonfirst-root fixture through the public factory.
+  新增 M2 实体消费层：物理源/variant 类型化、显式逻辑到物理绑定、局部实体 reset 事务和逐 variant 状态/控制默认值。注册 primitive task 验证同一可 pickle EnvFactory 在 MuJoCo/IsaacSim 的行为。标准与 ROCm 依赖配置要求已发布的 `unisim-core>=1.5.0` 包，不使用 Git source 覆盖。回归测试使用统一 worker，并通过公共 factory 构造非首 free root 场景。
+- Remove quadratic selected-row lookup and unused full-batch snapshots from entity reset staging. Reject mapped logical roots that point at descendant bodies, keeping reads, defaults and writes aligned. A bounded A/B script records separate row-index and sparse-field effects without claiming simulation throughput.
+  删除实体 reset 暂存中的平方级行查找和无用全批状态快照。拒绝指向后代 body 的逻辑 root 绑定，使读取、默认值和写入保持一致。限定规模的 A/B 脚本分别记录行索引与稀疏字段的影响，不宣称仿真吞吐提升。
+
 - Added IsaacGym fixed-variant protocol and real-runtime coverage. The
   deterministic worker mock validates and echoes the construction-time variant
   assignment, while the external Preview-4 slow lane realizes per-env actor

docs/sphinx/source/en/4-developer_guide/2-contracts/2-backend_contract.md

@@ -34,8 +34,22 @@ Optional capabilities are explicit:
 - Asset/XML/model metadata access belongs to cold paths such as scene
   materialization, backend init, or cache creation.
 
+## Physical entities and selected reset
+
+The M2 consumer in issue #1599 uses UniSim physical entity declarations separately from UniLab logical selectors. `SceneCfg` materializes typed entity/variant values, calls the parent contract validation, and the asset factory collects physical and catalog source paths. `EntityCfg.physical_entity` binds a logical facade explicitly; `primary_entity` selects the scene's primary root without encoding a task name in the backend.
+
+A mapped logical root must name its physical entity's declared root exactly. Binding a descendant body as the root is rejected during initialization, because root queries, defaults and reset writes must reference the same object. Reset staging maps selected rows in linear time and stores only requested fields; a current-state snapshot is needed only to fill missing columns when merging different joint position/velocity selections. The transaction still validates before its single public backend commit.
+
+The existing `ResetStateTransaction` stages one public `SceneResetRequest` for mapped scenes. Missing fields and unselected entities/environments remain unchanged. Per-environment defaults come from `get_entity_default_state`, and `restore_default_controls` restores keyframe controls in the same commit; controls need not equal joint positions. No engine-private tensors or asset parsing enter manager terms. Scalar hinge/slide joints and a common selected environment set per transaction are the current consumer boundary; unsupported mixed DR/mocap or row patterns fail explicitly.
+
+`tests/envs/test_multi_entity_consumer.py` registers one primitive task with the same pickleable EnvFactory for MuJoCo and IsaacSim. It checks observation/action dimensions, passive joints, selected resets, variants and a kinematic mirror. The native IsaacSim cases require `UNILAB_TEST_M2_ISAACSIM=1`. The consumer requires released `unisim-core>=1.5.0`; normal and ROCm lock profiles resolve the PyPI package without a Git source override. `UNILAB_LOCAL_UNISIM` remains an explicit alternative for local development. See [UniSim roadmap #108](https://github.com/unilabsim/unisim/issues/108) for implementation and verification scope.
+
 ## Evidence In Repo
 
+- Configuration and asset preparation: `src/unilab/base/scene.py`, `src/unilab/base/backend_factory.py`.
+- Public state/reset bindings: `src/unilab/base/entity.py`, `src/unilab/base/reset_state.py`.
+- Registered runtime tests: `tests/base/test_entity_scene_consumer.py`, `tests/envs/test_multi_entity_consumer.py`.
+
 - Backend interface and play capabilities: `unisim.backend.base`
 - Backend factory: `src/unilab/base/backend_factory.py`
 - MuJoCo backend: `unisim.backend.mujoco.backend`

docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/2-backend_contract.md

@@ -29,8 +29,22 @@ Backend 差异是契约边界，而不是脚本层面的特殊处理。play/rend
 - 资源/XML/模型元数据的访问属于冷路径，例如场景 materialization、backend init
   或 cache 创建。
 
+## 物理实体与局部 reset
+
+Issue #1599 的 M2 消费层将 UniSim 物理实体声明与 UniLab 逻辑 selector 分开。`SceneCfg` 物化实体/variant 类型并调用父契约校验，asset factory 收集物理源和 catalog 源路径。`EntityCfg.physical_entity` 显式绑定逻辑 facade；`primary_entity` 选择场景主根，不把任务名称写进 backend。
+
+映射的逻辑 root 必须精确指向物理实体声明的 root。初始化时拒绝将后代 body 绑定为 root，确保 root 查询、默认值和 reset 写入引用同一对象。Reset 暂存以线性时间映射选中行，并只保存请求的字段；只有合并不同 joint position/velocity 选择、需要补齐未写列时才读取当前状态快照。事务仍先完整校验，再调用一次公共 backend 提交。
+
+既有 `ResetStateTransaction` 为 mapped scene 暂存一次公共 `SceneResetRequest`。缺失字段、未选实体和环境保持不变。逐环境默认值来自 `get_entity_default_state`，`restore_default_controls` 在同次提交中恢复 keyframe control，控制值不必等于关节位置。Manager term 不接触引擎私有 tensor 或资产解析。当前消费边界为标量 hinge/slide 和一次事务共用选中环境集合；不支持的 DR/mocap 混写或行模式明确拒绝。
+
+`tests/envs/test_multi_entity_consumer.py` 为 MuJoCo 和 IsaacSim 注册同一个 primitive task，并使用同一个可 pickle 的 EnvFactory。测试检查观测/动作维度、被动关节、局部 reset、variants 和 kinematic mirror。原生 IsaacSim case 通过 `UNILAB_TEST_M2_ISAACSIM=1` 启用。消费层要求已发布的 `unisim-core>=1.5.0`；标准与 ROCm 锁文件均解析 PyPI 包，不使用 Git source 覆盖。`UNILAB_LOCAL_UNISIM` 仍是显式本地开发替代方案。实现和验证边界见 [UniSim roadmap #108](https://github.com/unilabsim/unisim/issues/108)。
+
 ## 仓库中的证据
 
+- 配置与资产准备：`src/unilab/base/scene.py`、`src/unilab/base/backend_factory.py`。
+- 公共状态/reset 绑定：`src/unilab/base/entity.py`、`src/unilab/base/reset_state.py`。
+- 注册 runtime 测试：`tests/base/test_entity_scene_consumer.py`、`tests/envs/test_multi_entity_consumer.py`。
+
 - Backend 接口与 play 能力：`unisim.backend.base`
 - Backend 工厂：`src/unilab/base/backend_factory.py`
 - MuJoCo backend：`unisim.backend.mujoco.backend`

pyproject.rocm.toml

@@ -24,10 +24,9 @@ requires-python = ">=3.10,<3.14"
 dependencies = [
     "numpy",
     # Physics implementations are provided by the independently released
-    # unisim-core package. The 1.4.2 release carries the fixed model variant,
-    # per-world reset-default, per-env gravity, substep body-wrench, and
-    # IsaacGym fixed-variant contracts.
-    "unisim-core>=1.4.2",
+    # unisim-core package. Version 1.5.0 provides the M2 physical entity,
+    # immutable variant, layout and selected-reset contracts.
+    "unisim-core>=1.5.0",
     # RL algorithms and async runtimes live in the independently released
     # uni-rl package (distribution name ``unilab-rl``); see pyproject.toml.
     "unilab-rl==1.2.1",

pyproject.toml

@@ -39,10 +39,9 @@ requires-python = ">=3.10,<3.14"
 dependencies = [
     "numpy",
     # Physics implementations are provided by the independently released
-    # unisim-core package. The 1.4.2 release carries the fixed model variant,
-    # per-world reset-default, per-env gravity, substep body-wrench, and
-    # IsaacGym fixed-variant contracts.
-    "unisim-core>=1.4.2",
+    # unisim-core package. Version 1.5.0 provides the M2 physical entity,
+    # immutable variant, layout and selected-reset contracts.
+    "unisim-core>=1.5.0",
     # RL algorithms and async runtimes (PPO/APPO/SAC/TD3 runners,
     # collectors, IPC, logging) live in the independently released uni-rl
     # package (distribution name ``unilab-rl``), consumed via the injected
@@ -170,7 +169,7 @@ viser = ["viser>=1.0.26", "trimesh>=3.21.7"]
 # required-environments; elsewhere the extra is empty and the CLI reports a
 # targeted runtime diagnostic.
 superdex = [
-    "unisim-core[superdex]>=1.4.2 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "unisim-core[superdex]>=1.5.0 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
 ]
 
 [dependency-groups]

scripts/benchmark/physics/m2_reset_ablation.py

@@ -0,0 +1,153 @@
+"""CPU host-fixture ablation of M2 reset staging; no physics or IPC timing.
+
+Run from a source checkout with its development dependencies installed. The
+baseline is loaded from local Git history, not from a temporary source file.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import platform
+import statistics
+import subprocess
+import sys
+import time
+import types
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[3]
+SOURCE_PATH = "src/unilab/base/reset_state.py"
+DEFAULT_BASELINE = "044a11ff"
+
+
+def _git(*args: str) -> str:
+    return subprocess.check_output(["git", *args], cwd=ROOT, text=True).strip()
+
+
+def run(baseline_revision: str, *, warmup: int, samples: int) -> dict:
+    """Compare old, row-map-only, and current sparse staging with equal outputs."""
+    if str(ROOT) not in sys.path:
+        sys.path.insert(0, str(ROOT))
+    import numpy as np
+    from tests.base.test_entity_reset_staging import fixture
+
+    from unilab.base.reset_state import ResetStateTransaction
+
+    baseline_sha = _git("rev-parse", "--verify", baseline_revision + "^{commit}")
+    source = _git("show", baseline_sha + ":" + SOURCE_PATH)
+    expression = "order = [request.env_ids.index(i) for i in rows]"
+    if source.count(expression) != 1:
+        raise ValueError("baseline must contain exactly one known quadratic row-order expression")
+    index_source = source.replace(
+        expression,
+        "incoming_rows = {value: index for index, value in enumerate(request.env_ids)}\n"
+        "        order = [incoming_rows[i] for i in rows]",
+        1,
+    )
+    paths = []
+    for name, code in (("A_git_baseline", source), ("B_row_index_only", index_source)):
+        module = types.ModuleType("_m2_ablation_" + name)
+        exec(compile(code, f"git:{baseline_sha}:{SOURCE_PATH}:{name}", "exec"), module.__dict__)
+        paths.append((name, module.ResetStateTransaction))
+    paths.append(("C_row_index_and_sparse_fields", ResetStateTransaction))
+    results = []
+    for count, selected in ((64, 1), (4096, 1), (4096, 1024), (4096, 4096)):
+        ids = np.arange(selected - 1, -1, -1)
+        pose = np.tile([0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 0.0], (selected, 1))
+        for kind in ("pose", "defaults"):
+            reference = None
+            for name, transaction_type in paths:
+                seed, reads, commits = fixture(count)
+                # The test fixture owns this simple stand-in; it is not a native backend.
+                transaction = transaction_type(seed._backend, scene_layout=seed.scene_layout)
+                times = []
+                for iteration in range(warmup + samples):
+                    reads.clear()
+                    commits.clear()
+                    start = time.perf_counter_ns()
+                    with transaction.scoped(ids):
+                        if kind == "pose":
+                            transaction.write_entity_state(
+                                "object", ids, term_name="ablation", root_pose=pose
+                            )
+                        else:
+                            transaction.reset_to_default(ids, term_name="ablation")
+                    elapsed = (time.perf_counter_ns() - start) / 1e6
+                    if iteration >= warmup:
+                        times.append(elapsed)
+                    if len(commits) != 1:
+                        raise AssertionError("every path must submit exactly one request")
+                    request = commits[0]
+                    patch = request.patches[0]
+                    signature = {
+                        key: None if getattr(patch, key) is None else getattr(patch, key).tolist()
+                        for key in (
+                            "root_pose",
+                            "root_velocity",
+                            "joint_positions",
+                            "joint_velocities",
+                        )
+                    }
+                    signature.update(
+                        joint_names=patch.joint_names,
+                        env_ids=request.env_ids,
+                        restore_default_controls=request.restore_default_controls,
+                    )
+                    if reference is None:
+                        reference = signature
+                    if signature != reference:
+                        raise AssertionError(
+                            f"semantic mismatch: {count}, {selected}, {kind}, {name}"
+                        )
+                results.append(
+                    {
+                        "N": count,
+                        "selected": selected,
+                        "kind": kind,
+                        "path": name,
+                        "median_ms": statistics.median(times),
+                        "min_ms": min(times),
+                        "max_ms": max(times),
+                        "current_snapshot_calls": len(reads),
+                        # Fixture's four float64 state arrays contain 7+6+2+2 values per env.
+                        # This counts getter-return bytes, not all Python/NumPy allocation.
+                        "current_snapshot_bytes": len(reads) * count * (7 + 6 + 2 + 2) * 8,
+                        "commit_calls": len(commits),
+                        "signature_matches_baseline": True,
+                    }
+                )
+    return {
+        "scope": "CPU host fixture only; no engine, IPC, worker, or training timing",
+        "baseline_sha": baseline_sha,
+        "current_head": _git("rev-parse", "HEAD"),
+        "working_tree_status": _git("status", "--short"),
+        "python": sys.version,
+        "platform": platform.platform(),
+        "numpy": np.__version__,
+        "warmup": warmup,
+        "samples": samples,
+        "results": results,
+    }
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--baseline", default=DEFAULT_BASELINE, help="local Git revision for path A"
+    )
+    parser.add_argument("--output", type=Path, required=True, help="JSON evidence destination")
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--samples", type=int, default=9)
+    args = parser.parse_args(argv)
+    if args.warmup < 0 or args.samples <= 0:
+        parser.error("warmup must be nonnegative and samples must be positive")
+    report = run(args.baseline, warmup=args.warmup, samples=args.samples)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
+    print(f"Saved {len(report['results'])} host-fixture measurements to {args.output}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/unilab/base/backend_factory.py

@@ -85,9 +85,13 @@ def create_backend(
         kwargs.pop("superdex_execution_mode", None)
         kwargs.pop("superdex_effort_limits", None)
         kwargs.pop("superdex_allow_contact_approximation", None)
-    ensure_robot_assets_for_paths(
-        [scene.model_file, scene.visual_model_file, *scene.fragment_files]
+    paths = [scene.model_file, scene.visual_model_file, *scene.fragment_files]
+    paths.extend(
+        entity.source.model_file for entity in scene.entity_assets if entity.source is not None
     )
+    if scene.entity_variant is not None:
+        paths.extend(source.model_file for source in scene.entity_variant.plan.variants)
+    ensure_robot_assets_for_paths(list(dict.fromkeys(paths)))
     if backend_type == "drake":
         # unisim-core 1.4.2 dropped the Drake-branch filtering of MuJoCo
         # root-body options; Drake derives root state from its own plant and

src/unilab/base/config_materialization.py

@@ -218,6 +218,11 @@ def _prepare_dataclass(
     hints = _hints(target)
     prepared: dict[str, Any] = {HYDRA_TARGET_KEY: reference}
     for name, raw_value in values.items():
+        from unilab.base.scene import SceneCfg
+
+        if issubclass(target, SceneCfg) and name in ("entity_assets", "entity_variant"):
+            prepared[name] = _plain(raw_value)
+            continue
         field = fields[name]
         field_path = f"{path}.{name}"
         annotation = hints.get(name, field.type)
@@ -365,6 +370,11 @@ def apply_cfg_overrides(
         if not isinstance(key, str) or not hasattr(target_obj, key):
             raise ValueError(f"Config class '{type(target_obj).__name__}' has no attribute '{key}'")
         value = _plain(raw_value)
+        from unilab.base.scene import SceneCfg
+
+        if isinstance(target_obj, SceneCfg) and key in ("entity_assets", "entity_variant"):
+            setattr(target_obj, key, value)
+            continue
         existing = getattr(target_obj, key)
         annotation = hints.get(key, fields[key].type if key in fields else Any)
         policy = _policy(target_obj, key)
@@ -408,5 +418,10 @@ def apply_cfg_overrides(
                 continue
         setattr(target_obj, key, _prepare_value(value, annotation=Any, path=path))
 
+    from unilab.base.scene import SceneCfg
+
+    if isinstance(target_obj, SceneCfg):
+        target_obj.__post_init__()
+
 
 __all__ = ["HYDRA_TARGET_KEY", "apply_cfg_overrides"]