diff --git a/docs/sphinx/source/changelog.md b/docs/sphinx/source/changelog.md
index aa1506530..a0e4c4894 100644
--- a/docs/sphinx/source/changelog.md
+++ b/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
diff --git a/docs/sphinx/source/en/4-developer_guide/2-contracts/2-backend_contract.md b/docs/sphinx/source/en/4-developer_guide/2-contracts/2-backend_contract.md
index 6bb442429..7800055fc 100644
--- a/docs/sphinx/source/en/4-developer_guide/2-contracts/2-backend_contract.md
+++ b/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`
diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/2-backend_contract.md b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/2-backend_contract.md
index 346b8be9d..401b4b59b 100644
--- a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/2-backend_contract.md
+++ b/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`
diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml
index 93770addb..b987516f9 100644
--- a/pyproject.rocm.toml
+++ b/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",
diff --git a/pyproject.toml b/pyproject.toml
index 5637d1bfe..6e061b4d0 100644
--- a/pyproject.toml
+++ b/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]
diff --git a/scripts/benchmark/physics/m2_reset_ablation.py b/scripts/benchmark/physics/m2_reset_ablation.py
new file mode 100644
index 000000000..fa29ad99a
--- /dev/null
+++ b/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())
diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py
index f4ffa1676..b4718e606 100644
--- a/src/unilab/base/backend_factory.py
+++ b/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
diff --git a/src/unilab/base/config_materialization.py b/src/unilab/base/config_materialization.py
index 4a070a48b..3480b4a08 100644
--- a/src/unilab/base/config_materialization.py
+++ b/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"]
diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py
index 162ea65c6..bc3ababc2 100644
--- a/src/unilab/base/entity.py
+++ b/src/unilab/base/entity.py
@@ -48,6 +48,7 @@ class EntityCfg:
geom_names: NamesCfg = None
site_names: NamesCfg = None
actuator_names: NamesCfg = None
+ physical_entity: str | None = None
def _normalize_names(entity_name: str, kind: str, names: NamesCfg) -> tuple[str, ...] | None:
@@ -587,12 +588,44 @@ def __init__(
self._joint_model_dof_ids: np.ndarray | None = None
self._motion_body_ids: np.ndarray | None = None
self._mocap_body_name: str | None = None
+ self._physical_entity: str | None = None
+ self._entity_defaults: dict[str, np.ndarray] | None = None
+ if reset_state is not None and reset_state.scene_layout is not None:
+ layout = reset_state.scene_layout
+ physical = cfg.physical_entity
+ if physical is None and cfg.root_body_name is not None and "/" in cfg.root_body_name:
+ physical = cfg.root_body_name.split("/", 1)[0]
+ if physical is None:
+ raise ValueError(f"Entity '{name}' requires physical_entity in a composed scene")
+ owner = layout.get_entity(physical)
+ expected_root = physical + "/" + owner.root_body
+ if cfg.root_body_name is not None and cfg.root_body_name != expected_root:
+ raise ValueError(
+ f"Entity '{name}' root_body_name must name physical root {expected_root!r}; "
+ f"got {cfg.root_body_name!r}"
+ )
+ if any(j.kind not in ("hinge", "slide") for j in owner.joints):
+ raise NotImplementedError("UniLab entity consumer currently supports scalar joints")
+ self._physical_entity = physical
+ self._entity_defaults = dict(backend.get_entity_default_state(physical))
self._joint_names = _normalize_names(name, "joint", cfg.joint_names)
self._body_names = _normalize_names(name, "body", cfg.body_names)
self._geom_names = _normalize_names(name, "geom", cfg.geom_names)
self._site_names = _normalize_names(name, "site", cfg.site_names)
self._actuator_names = _normalize_names(name, "actuator", cfg.actuator_names)
+ if self._physical_entity is not None:
+ prefix = self._physical_entity + "/"
+ selected = [
+ value
+ for values in (self._joint_names, self._body_names, self._actuator_names)
+ if values
+ for value in values
+ ]
+ if cfg.root_body_name:
+ selected.append(cfg.root_body_name)
+ if any(not value.startswith(prefix) for value in selected):
+ raise ValueError("mapped logical selectors must belong to their physical entity")
root_body_ids = None
if cfg.root_body_name is not None:
@@ -760,7 +793,7 @@ def _resolve_enumerated_ids(
f"backend '{self._backend_type}'; available={list(all_names)}"
)
return _readonly_ids(
- [ids_by_name[value] for value in names],
+ np.asarray([ids_by_name[value] for value in names], dtype=np.int32),
expected=len(names),
label=f"Entity '{self.name}' {capability}",
)
@@ -868,6 +901,8 @@ def _materialize_default_joint_pos(
) -> np.ndarray | None:
if joint_pos_ids is None:
return None
+ if self._entity_defaults is not None:
+ return self._selected_entity_default_joints("joint_positions")
current = self._read_state("joint position state", backend.get_dof_pos)
if default_qpos is None:
defaults = self._read_state("default joint position", backend.get_default_dof_pos)
@@ -911,6 +946,10 @@ def _materialize_soft_joint_pos_limits(
) -> np.ndarray | None:
if joint_pos_ids is None:
return None
+ if self._physical_entity is not None:
+ result = np.asarray(backend.get_joint_range(names=self._joint_names)).copy()
+ result.setflags(write=False)
+ return result
try:
raw_ranges = backend.get_joint_range()
except (AttributeError, NotImplementedError) as exc:
@@ -941,6 +980,11 @@ def _materialize_root_state(
) -> tuple[BackendRootStateLayout | None, np.ndarray | None, str | None]:
if root_body_name is None:
return None, None, "root_body_name was not declared in EntityCfg"
+ if self._entity_defaults is not None:
+ defaults = self._entity_defaults
+ value = np.concatenate((defaults["root_pose"], defaults["root_velocity"]), axis=1)
+ value.setflags(write=False)
+ return None, value, None
try:
layout = backend.get_root_state_layout(root_body_name)
except (AttributeError, NotImplementedError) as exc:
@@ -1011,6 +1055,8 @@ def _materialize_default_joint_vel(
) -> np.ndarray | None:
if joint_vel_ids is None:
return None
+ if self._entity_defaults is not None:
+ return self._selected_entity_default_joints("joint_velocities")
current = self._read_state("joint velocity state", backend.get_dof_vel)
materialized = np.zeros(
(backend.num_envs, len(joint_vel_ids)),
@@ -1019,6 +1065,16 @@ def _materialize_default_joint_vel(
materialized.setflags(write=False)
return materialized
+ def _selected_entity_default_joints(self, field: str) -> np.ndarray:
+ assert self._reset_state is not None and self._reset_state.scene_layout is not None
+ assert self._physical_entity is not None and self._entity_defaults is not None
+ owner = self._reset_state.scene_layout.get_entity(self._physical_entity)
+ names = [joint.name for joint in owner.joints]
+ selected = [names.index(name.split("/", 1)[1]) for name in self._joint_names or ()]
+ result = self._entity_defaults[field][:, selected].copy()
+ result.setflags(write=False)
+ return result
+
def _materialize_gravity_vector(
self, backend: SimBackend, root_body_ids: np.ndarray | None
) -> np.ndarray | None:
@@ -1309,6 +1365,13 @@ def write_root_state_to_sim(
env_ids: np.ndarray | slice | None = None,
) -> None:
"""Stage a 13-D world-frame root state in the active reset transaction."""
+ if self._physical_entity is not None:
+ if root_state.ndim != 2 or root_state.shape[1] != 13:
+ raise ValueError("root_state must have 13 columns")
+ self._stage_entity_write(
+ env_ids, root_pose=root_state[:, :7], root_velocity=root_state[:, 7:]
+ )
+ return
reset_state, layout = self._require_root_state_write()
resolved_env_ids = self._normalize_reset_env_ids(env_ids)
reset_state.write_root_state(
@@ -2084,6 +2147,9 @@ def write_root_link_pose_to_sim(
env_ids: np.ndarray | slice | None = None,
) -> None:
"""Stage world position and wxyz root orientation during reset."""
+ if self._physical_entity is not None:
+ self._stage_entity_write(env_ids, root_pose=root_pose)
+ return
reset_state, layout = self._require_root_state_write()
resolved_env_ids = self._normalize_reset_env_ids(env_ids)
reset_state.write_root_pose(
@@ -2099,6 +2165,9 @@ def write_root_link_velocity_to_sim(
env_ids: np.ndarray | slice | None = None,
) -> None:
"""Stage world linear/angular root velocity during reset."""
+ if self._physical_entity is not None:
+ self._stage_entity_write(env_ids, root_velocity=root_velocity)
+ return
reset_state, layout = self._require_root_state_write()
resolved_env_ids = self._normalize_reset_env_ids(env_ids)
reset_state.write_root_velocity(
@@ -2119,6 +2188,11 @@ def read_reset_root_pose(
term has written yet), so a later reset term can build on an earlier
term's root placement.
"""
+ if self._physical_entity is not None:
+ assert self._reset_state is not None
+ return self._reset_state.read_entity_root_pose(
+ self._physical_entity, self._normalize_reset_env_ids(env_ids)
+ )
reset_state, layout = self._require_root_state_write()
resolved_env_ids = self._normalize_reset_env_ids(env_ids)
return reset_state.read_root_pose(
@@ -2163,6 +2237,15 @@ def write_joint_state_to_sim(
capability="reset joint-state write",
)
resolved_env_ids = self._normalize_reset_env_ids(env_ids)
+ if self._physical_entity is not None:
+ names = tuple(self._joint_names[int(i)].split("/", 1)[1] for i in local_joint_ids)
+ self._stage_entity_write(
+ resolved_env_ids,
+ joint_positions=position,
+ joint_velocities=velocity,
+ joint_names=names,
+ )
+ return
self._materialize_reset_joint_indices()
assert self._reset_joint_qpos_ids is not None
assert self._reset_joint_qvel_ids is not None
@@ -2175,6 +2258,15 @@ def write_joint_state_to_sim(
term_name=f"{self.name}.write_joint_state_to_sim",
)
+ def _stage_entity_write(self, env_ids, **fields) -> None:
+ assert self._reset_state is not None and self._physical_entity is not None
+ self._reset_state.write_entity_state(
+ self._physical_entity,
+ self._normalize_reset_env_ids(env_ids),
+ term_name=f"{self.name}.entity_state",
+ **fields,
+ )
+
def _materialize_reset_joint_indices(self) -> None:
if self._reset_joint_qpos_ids is not None and self._reset_joint_qvel_ids is not None:
return
diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py
index d3c3814f1..4beec1a0a 100644
--- a/src/unilab/base/reset_state.py
+++ b/src/unilab/base/reset_state.py
@@ -30,6 +30,8 @@
RESET_TERM_KP,
ResetRandomizationPayload,
)
+from unisim.entities import EntityStatePatch, SceneResetRequest
+from unisim.scene_layout import CompiledSceneLayout
from unilab.utils.rotation import np_quat_apply_inverse
@@ -73,10 +75,34 @@ def __init__(
backend: SimBackend,
*,
default_qpos: np.ndarray | None = None,
+ scene_layout: CompiledSceneLayout | None = None,
) -> None:
self._backend = backend
self._num_envs = backend.num_envs
self._selected_default_qpos = default_qpos
+ self.scene_layout = scene_layout
+ if scene_layout is not None and any(
+ joint.kind not in ("hinge", "slide")
+ for entity in scene_layout.entities
+ for joint in entity.joints
+ ):
+ raise NotImplementedError("mapped manager reset currently supports scalar joints")
+ self._entity_values: dict[str, dict[str, np.ndarray]] = {}
+ self._entity_fields: dict[str, set[str]] = {}
+ self._entity_joints: dict[str, set[str]] = {}
+ self._entity_joint_fields: dict[str, dict[str, set[str]]] = {}
+ self._entity_layouts = (
+ {}
+ if scene_layout is None
+ else {entity.name: entity for entity in scene_layout.entities}
+ )
+ self._entity_joint_columns = {
+ name: {joint.name: index for index, joint in enumerate(entity.joints)}
+ for name, entity in self._entity_layouts.items()
+ }
+ self._entity_rows: tuple[int, ...] | None = None
+ self._entity_row_index: dict[int, int] = {}
+ self._restore_entity_controls = False
self._active = False
self._active_mask = np.zeros(self._num_envs, dtype=np.bool_)
self._dirty_mask = np.zeros(self._num_envs, dtype=np.bool_)
@@ -151,8 +177,124 @@ def begin(self, env_ids: np.ndarray) -> None:
self._requesting_terms.clear()
self._last_commit_had_writes = False
self._last_set_state_timing_ms = {}
+ self._entity_values.clear()
+ self._entity_fields.clear()
+ self._entity_joints.clear()
+ self._entity_joint_fields.clear()
+ self._entity_rows = None
+ self._entity_row_index.clear()
+ self._restore_entity_controls = False
self._active = True
+ def write_entity_state(
+ self,
+ entity: str,
+ env_ids: np.ndarray,
+ *,
+ term_name: str,
+ root_pose: np.ndarray | None = None,
+ root_velocity: np.ndarray | None = None,
+ joint_positions: np.ndarray | None = None,
+ joint_velocities: np.ndarray | None = None,
+ joint_names: tuple[str, ...] = (),
+ ) -> None:
+ """Stage one public physical-entity patch in the current manager transaction."""
+ self._require_active()
+ if self.scene_layout is None:
+ raise NotImplementedError("physical entity patches require a compiled scene layout")
+ ids = self._validate_ids(env_ids, capability=term_name)
+ if np.any(~self._active_mask[ids]):
+ raise ValueError("entity write is outside the active reset")
+ if not len(ids):
+ return
+ patch = EntityStatePatch(
+ entity, root_pose, root_velocity, joint_positions, joint_velocities, joint_names
+ )
+ request = SceneResetRequest(tuple(int(i) for i in ids), (patch,))
+ self.scene_layout.validate_reset(request, num_envs=self._num_envs)
+ rows = tuple(sorted(request.env_ids))
+ if self._entity_rows is not None and self._entity_rows != rows:
+ raise NotImplementedError(
+ "one entity transaction requires the same selected env rows for every patch"
+ )
+ self._entity_rows = rows
+ self._entity_row_index = {value: index for index, value in enumerate(rows)}
+ incoming_rows = {value: index for index, value in enumerate(request.env_ids)}
+ order = [incoming_rows[i] for i in rows]
+ owner = self._entity_layouts[entity]
+ if entity not in self._entity_values:
+ self._entity_values[entity] = {}
+ self._entity_fields[entity] = set()
+ self._entity_joints[entity] = set()
+ self._entity_joint_fields[entity] = {}
+ values = self._entity_values[entity]
+ selected_joints = joint_names or tuple(joint.name for joint in owner.joints)
+ for field in ("root_pose", "root_velocity", "joint_positions", "joint_velocities"):
+ incoming = getattr(patch, field)
+ if incoming is None:
+ continue
+ if field.startswith("joint"):
+ columns = [self._entity_joint_columns[entity][name] for name in selected_joints]
+ if field not in values:
+ values[field] = np.empty((len(rows), len(owner.joints)), dtype=np.float64)
+ values[field][:, columns] = incoming[order]
+ self._entity_joints[entity].update(selected_joints)
+ self._entity_joint_fields[entity].setdefault(field, set()).update(selected_joints)
+ else:
+ values[field] = incoming[order]
+ self._entity_fields[entity].add(field)
+ self._requesting_terms.add(term_name)
+
+ def read_entity_root_pose(self, entity: str, env_ids: np.ndarray) -> np.ndarray:
+ """Read staged root pose or a detached current public snapshot."""
+ self._require_active()
+ env_ids = self._validate_ids(env_ids, capability="read_entity_root_pose")
+ if np.any(~self._active_mask[env_ids]):
+ raise ValueError("entity state read is outside the active reset")
+ if "root_pose" in self._entity_values.get(entity, {}):
+ rows = [self._entity_row_index[int(i)] for i in env_ids]
+ return self._entity_values[entity]["root_pose"][rows].copy()
+ return self._backend.get_entity_state(entity)["root_pose"][env_ids].copy()
+
+ def _commit_entities(self) -> None:
+ assert self.scene_layout is not None
+ if self._entity_rows is None:
+ return
+ if np.any(self._dirty_mask) or any(np.any(mask) for mask in self._mocap_masks.values()):
+ raise NotImplementedError("cannot mix entity patches with legacy state/DR writes")
+ patches = []
+ for name, values in self._entity_values.items():
+ entity = self._entity_layouts[name]
+ joint_names = tuple(
+ j.name for j in entity.joints if j.name in self._entity_joints[name]
+ )
+ columns = [self._entity_joint_columns[name][n] for n in joint_names]
+ current = None
+ for field, written in self._entity_joint_fields[name].items():
+ missing = self._entity_joints[name] - written
+ if missing:
+ if current is None:
+ current = self._backend.get_entity_state(name)
+ absent = [self._entity_joint_columns[name][joint] for joint in missing]
+ values[field][:, absent] = current[field][np.ix_(self._entity_rows, absent)]
+ fields = {
+ field: values[field][:, columns] if field.startswith("joint") else values[field]
+ for field in self._entity_fields[name]
+ }
+ patches.append(EntityStatePatch(name, joint_names=joint_names, **fields))
+ request = SceneResetRequest(
+ self._entity_rows,
+ tuple(patches),
+ restore_default_controls=self._restore_entity_controls,
+ )
+ self.scene_layout.validate_reset(request, num_envs=self._num_envs)
+ started = time.perf_counter()
+ self._backend.reset_entities(request)
+ self._last_commit_had_writes = True
+ self._last_set_state_timing_ms = {
+ "dr_reset_set_state_ms": (time.perf_counter() - started) * 1000
+ }
+
def bind_geom_size_write(
self,
column_ids: np.ndarray,
@@ -728,6 +870,21 @@ def write_actuator_gains(
def reset_to_default(self, env_ids: np.ndarray, *, term_name: str) -> None:
"""Stage backend default qpos/qvel for a subset of the active reset."""
self._require_active()
+ if self.scene_layout is not None:
+ self._restore_entity_controls = True
+ for entity in self.scene_layout.entities:
+ defaults = self._backend.get_entity_default_state(entity.name, env_ids)
+ fields = {}
+ if entity.root_mode != "fixed":
+ fields["root_pose"] = defaults["root_pose"]
+ if entity.root_mode == "floating":
+ fields["root_velocity"] = defaults["root_velocity"]
+ if entity.joints:
+ fields["joint_positions"] = defaults["joint_positions"]
+ fields["joint_velocities"] = defaults["joint_velocities"]
+ if fields:
+ self.write_entity_state(entity.name, env_ids, term_name=term_name, **fields)
+ return
ids = self._validate_ids(env_ids, capability="reset_to_default")
outside = ids[~self._active_mask[ids]]
if outside.size:
@@ -954,6 +1111,16 @@ def write_mocap_pose(
def commit(self) -> dict | None:
"""Commit all staged rows through one public backend call."""
self._require_active()
+ if self.scene_layout is not None:
+ try:
+ self._commit_entities()
+ if np.any(self._dirty_mask):
+ raise NotImplementedError(
+ "mapped scene DR writes need a public entity transaction"
+ )
+ return None
+ finally:
+ self._finish()
dirty_ids = np.flatnonzero(self._dirty_mask).astype(np.int32, copy=False)
mocap_dirty = any(np.any(mask) for mask in self._mocap_masks.values())
self._last_commit_had_writes = bool(dirty_ids.size) or mocap_dirty
diff --git a/src/unilab/base/scene.py b/src/unilab/base/scene.py
index 10668cec7..fbd57ef57 100644
--- a/src/unilab/base/scene.py
+++ b/src/unilab/base/scene.py
@@ -5,6 +5,9 @@
from collections.abc import Mapping
from dataclasses import dataclass, field
+import numpy as np
+from unisim.dr.types import FixedVariantLayout, FixedVariantPlan, ModelSourceDescriptor
+from unisim.entities import EntityInitialState, EntityVariantBinding, SceneEntitySpec
from unisim.scene import (
SceneCfg as _UniSimSceneCfg,
)
@@ -26,12 +29,58 @@ class SceneCfg(_UniSimSceneCfg):
manager facade's :class:`EntityCfg` records on the cold configuration path.
"""
- model_file: str
+ model_file: str = ""
fragment_files: list[str] = field(default_factory=list)
terrain: TerrainSceneCfg | None = None
entities: dict[str, object] = field(default_factory=dict)
+ entity_assets: tuple[SceneEntitySpec, ...] = ()
+ entity_variant: EntityVariantBinding | None = None
+ primary_entity: str | None = None
+
+ def materialize_entities(self) -> None:
+ """Convert task/Hydra physical records without importing a physics SDK."""
+ physical = []
+ for raw in self.entity_assets:
+ if isinstance(raw, SceneEntitySpec):
+ physical.append(raw)
+ continue
+ if not isinstance(raw, Mapping):
+ raise TypeError("entity_assets entries must be SceneEntitySpec or mappings")
+ values = dict(raw)
+ source = values.get("source")
+ if isinstance(source, Mapping):
+ values["source"] = ModelSourceDescriptor(**source)
+ state = values.get("initial_state")
+ if isinstance(state, Mapping):
+ values["initial_state"] = EntityInitialState(
+ **{key: tuple(value) for key, value in state.items()}
+ )
+ physical.append(SceneEntitySpec(**values))
+ self.entity_assets = tuple(physical)
+ if isinstance(self.entity_variant, Mapping):
+ binding = dict(self.entity_variant)
+ plan = binding.get("plan")
+ if isinstance(plan, Mapping):
+ data = dict(plan)
+ data["assignment"] = np.asarray(data["assignment"])
+ data["variants"] = tuple(
+ ModelSourceDescriptor(**source) if isinstance(source, Mapping) else source
+ for source in data["variants"]
+ )
+ if "layout" in data:
+ data["layout"] = FixedVariantLayout(data["layout"])
+ binding["plan"] = FixedVariantPlan(**data)
+ self.entity_variant = EntityVariantBinding(**binding)
+ if (
+ self.entity_assets
+ and self.primary_entity is not None
+ and self.primary_entity not in {entity.name for entity in self.entity_assets}
+ ):
+ raise ValueError("primary_entity must name a declared physical entity")
+ self.validate_composition()
def __post_init__(self) -> None:
+ self.materialize_entities()
materialized: dict[str, object] = {}
for name, value in self.entities.items():
if isinstance(value, EntityCfg):
@@ -41,6 +90,7 @@ def __post_init__(self) -> None:
else:
materialized[name] = value
self.entities = materialized
+ super().__post_init__()
__all__ = [
diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py
index bc2c53ff5..3a36d1b51 100644
--- a/src/unilab/envs/manager_based_rl_env.py
+++ b/src/unilab/envs/manager_based_rl_env.py
@@ -151,11 +151,40 @@ def validate(self) -> None:
"ManagerBasedRlEnvCfg scene must be a SceneCfg instance, "
f"got {type(self.scene).__name__}"
)
+ self.scene.__post_init__()
def _resolve_backend_entity_contract(cfg: ManagerBasedRlEnvCfg) -> tuple[str, bool]:
"""Resolve task-independent backend inputs from declared scene entities."""
assert cfg.scene is not None
+ if cfg.scene.entity_assets:
+ names = {entity.name: entity for entity in cfg.scene.entity_assets}
+ primary = cfg.scene.primary_entity
+ if primary is None:
+ controlled = [
+ name
+ for name, selector in cfg.scene.entities.items()
+ if isinstance(selector, EntityCfg) and selector.actuator_names
+ ]
+ if len(controlled) != 1:
+ raise ValueError("composed manager scene requires explicit primary_entity")
+ selector = cfg.scene.entities[controlled[0]]
+ assert isinstance(selector, EntityCfg)
+ primary = selector.physical_entity
+ if primary is None and selector.root_body_name:
+ primary = selector.root_body_name.split("/", 1)[0]
+ if primary is None or primary not in names:
+ raise ValueError("primary entity must name a physical scene entity")
+ root_names = [
+ selector.root_body_name
+ for selector in cfg.scene.entities.values()
+ if isinstance(selector, EntityCfg)
+ and selector.root_body_name
+ and selector.root_body_name.startswith(primary + "/")
+ ]
+ if not root_names:
+ raise ValueError("primary_entity needs an entity-qualified logical root selector")
+ return root_names[0], True
root_entities: list[tuple[str, str]] = []
body_state_requested = False
for entity_name, entity_cfg in cfg.scene.entities.items():
@@ -186,14 +215,14 @@ def _resolve_backend_entity_contract(cfg: ManagerBasedRlEnvCfg) -> tuple[str, bo
"ManagerBasedRlEnv factory requires at least one scene entity with an explicit "
"root_body_name"
)
- primary = next((item for item in root_entities if item[0] == "robot"), None)
- if primary is None and len(root_entities) != 1:
+ primary_root = next((item for item in root_entities if item[0] == "robot"), None)
+ if primary_root is None and len(root_entities) != 1:
declared = [name for name, _ in root_entities]
raise ValueError(
"ManagerBasedRlEnv factory requires a conventional 'robot' root entity when "
f"multiple floating entities are declared; found {declared}"
)
- return (primary or root_entities[0])[1], body_state_requested
+ return (primary_root or root_entities[0])[1], body_state_requested
class ManagerBasedRlEnv(NpEnv):
@@ -235,7 +264,13 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int
assert cfg.scene is not None
default_qpos = resolve_scene_default_qpos(cfg.scene, backend)
self._control = np.zeros((num_envs, backend.num_actuators), dtype=get_global_dtype())
- self._reset_state = ResetStateTransaction(backend, default_qpos=default_qpos)
+ if cfg.scene.entity_assets:
+ self._control[:] = backend.get_state("ctrl")["ctrl"]
+ self._reset_state = ResetStateTransaction(
+ backend,
+ default_qpos=default_qpos,
+ scene_layout=backend.get_scene_layout() if cfg.scene.entity_assets else None,
+ )
self.scene = EntityScene.from_scene_cfg(
cfg.scene,
backend,
@@ -606,7 +641,10 @@ def reset(
log.update(manager.reset(ids))
self.episode_length_buf[ids] = 0
- self._control[ids] = 0.0
+ if self._reset_state.scene_layout is not None:
+ self._control[ids] = self._backend.get_state("ctrl")["ctrl"][ids]
+ else:
+ self._control[ids] = 0.0
self._manual_reset_pending[ids] = False
if self._state is not None:
self._state.info["steps"][ids] = 0
diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py
index f7ec1b44d..e22cdf4ef 100644
--- a/tests/base/test_backend_conformance.py
+++ b/tests/base/test_backend_conformance.py
@@ -399,16 +399,13 @@ def test_root_qvel_body_angular_contract_reads_back_world_velocity(backend_type:
)
-def test_mujoco_root_layout_resolves_a_nonfirst_free_joint() -> None:
- import mujoco
-
+def test_mujoco_root_layout_resolves_a_nonfirst_free_joint(tmp_path: Path) -> None:
pytest.importorskip(
"unisim.backend.mujoco.backend",
reason="unisim-core MuJoCo adapter (mjbatch build) not available",
)
- from unisim.backend.mujoco.backend import MuJoCoBackend
-
- model = mujoco.MjModel.from_xml_string(
+ source = tmp_path / "nonfirst-root.xml"
+ source.write_text(
"""
@@ -424,14 +421,17 @@ def test_mujoco_root_layout_resolves_a_nonfirst_free_joint() -> None:
"""
)
- backend = object.__new__(MuJoCoBackend)
- backend._model = model
-
- layout = backend.get_root_state_layout("floating")
- assert layout.qpos_indices == tuple(range(1, 8))
- assert layout.qvel_indices == tuple(range(1, 7))
- with pytest.raises(NotImplementedError, match="hinged.*exactly one free joint"):
- backend.get_root_state_layout("hinged")
+ backend = create_backend(
+ "mujoco", SceneCfg(model_file=str(source)), 1, SIM_DT, base_name="floating"
+ )
+ try:
+ layout = backend.get_root_state_layout("floating")
+ assert layout.qpos_indices == tuple(range(1, 8))
+ assert layout.qvel_indices == tuple(range(1, 7))
+ with pytest.raises(NotImplementedError, match="hinged.*exactly one free joint"):
+ backend.get_root_state_layout("hinged")
+ finally:
+ backend.close()
def test_drake_root_layout_is_explicitly_unsupported_without_runtime_metadata() -> None:
diff --git a/tests/base/test_backend_pre_step_control.py b/tests/base/test_backend_pre_step_control.py
index 1516db424..8500ee270 100644
--- a/tests/base/test_backend_pre_step_control.py
+++ b/tests/base/test_backend_pre_step_control.py
@@ -136,6 +136,10 @@ def _fake_mujoco_backend(pre_step_control_fn=None):
pool = _FakeMjBatch()
backend = object.__new__(MuJoCoBackend)
+ backend._entity_layout = None
+ backend._tracked_body_state_dirty = np.zeros(1, dtype=np.bool_)
+ backend._tracked_sensor_slices = {}
+ backend._refresh_pre_step_body_state = True
backend._pre_step_control_fn = pre_step_control_fn
backend._tracked_sensor_copyout_range = None
backend._num_envs = 1
diff --git a/tests/base/test_entity_reset_staging.py b/tests/base/test_entity_reset_staging.py
new file mode 100644
index 000000000..012956a79
--- /dev/null
+++ b/tests/base/test_entity_reset_staging.py
@@ -0,0 +1,157 @@
+"""Sparse reset staging must avoid unused snapshots without changing writes."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from unisim.scene_layout import CompiledSceneLayout, EntityLayout, JointLayout
+
+from unilab.base.reset_state import ResetStateTransaction
+
+
+def fixture(num_envs=5):
+ entity = EntityLayout(
+ "object",
+ "articulation",
+ "floating",
+ "base",
+ ("base", "a", "b"),
+ (0, 1, 2),
+ (None, "base", "a"),
+ (JointLayout("a", "hinge", (7,), (6,), "a"), JointLayout("b", "hinge", (8,), (7,), "b")),
+ (),
+ (),
+ (),
+ tuple(range(7)),
+ tuple(range(6)),
+ )
+ layout = CompiledSceneLayout((entity,), 9, 8, 0, 3)
+ state = {
+ "root_pose": np.tile([0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0], (num_envs, 1)),
+ "root_velocity": np.zeros((num_envs, 6)),
+ "joint_positions": np.tile([0.2, 0.4], (num_envs, 1)),
+ "joint_velocities": np.tile([0.6, 0.8], (num_envs, 1)),
+ }
+ reads, commits = [], []
+
+ def read(name):
+ reads.append(name)
+ return {key: values.copy() for key, values in state.items()}
+
+ backend = SimpleNamespace(
+ num_envs=num_envs,
+ get_entity_state=read,
+ reset_entities=commits.append,
+ get_entity_default_state=lambda name, ids: {
+ key: values[ids].copy() for key, values in state.items()
+ },
+ )
+ return ResetStateTransaction(backend, scene_layout=layout), reads, commits
+
+
+@pytest.mark.parametrize("default", [False, True])
+def test_complete_fields_do_not_read_current_entity_state(default):
+ transaction, reads, commits = fixture()
+ ids = np.array([4, 1])
+ with transaction.scoped(ids):
+ if default:
+ transaction.reset_to_default(ids, term_name="defaults")
+ else:
+ transaction.write_entity_state(
+ "object",
+ ids,
+ term_name="pose",
+ root_pose=np.array([[4.0, 0, 1, 1, 0, 0, 0], [1.0, 0, 1, 1, 0, 0, 0]]),
+ )
+ assert reads == [] and len(commits) == 1
+ assert commits[0].env_ids == (1, 4)
+ if not default:
+ np.testing.assert_array_equal(commits[0].patches[0].root_pose[:, 0], [1, 4])
+
+
+def test_mixed_joint_field_subsets_read_once_only_for_missing_merge_columns():
+ transaction, reads, commits = fixture()
+ ids = np.array([4, 1])
+ with transaction.scoped(ids):
+ transaction.write_entity_state(
+ "object", ids, term_name="a", joint_names=("a",), joint_positions=np.array([[1], [2]])
+ )
+ transaction.write_entity_state(
+ "object",
+ ids[::-1],
+ term_name="b",
+ joint_names=("b",),
+ joint_velocities=np.array([[3.5], [4.5]]),
+ )
+ assert reads == []
+ assert reads == ["object"]
+ patch = commits[0].patches[0]
+ assert patch.joint_names == ("a", "b")
+ np.testing.assert_allclose(patch.joint_positions, [[2, 0.4], [1, 0.4]])
+ np.testing.assert_allclose(patch.joint_velocities, [[0.6, 3.5], [0.6, 4.5]])
+
+
+def test_overlapping_joint_writes_preserve_fractions_and_last_write_wins():
+ transaction, reads, commits = fixture()
+ with transaction.scoped(np.array([1])):
+ transaction.write_entity_state(
+ "object",
+ np.array([1]),
+ term_name="first",
+ joint_names=("a",),
+ joint_positions=np.array([[1]]),
+ )
+ transaction.write_entity_state(
+ "object",
+ np.array([1]),
+ term_name="second",
+ joint_names=("a",),
+ joint_positions=np.array([[0.25]]),
+ )
+ assert reads == []
+ np.testing.assert_allclose(commits[0].patches[0].joint_positions, [[0.25]])
+
+
+def test_late_invalid_patch_does_not_read_current_or_commit():
+ transaction, reads, commits = fixture()
+ with pytest.raises(ValueError), transaction.scoped(np.array([1])):
+ transaction.write_entity_state(
+ "object",
+ np.array([1]),
+ term_name="first",
+ joint_names=("a",),
+ joint_positions=np.array([[0.25]]),
+ )
+ transaction.write_entity_state(
+ "object", np.array([1]), term_name="invalid", root_pose=np.zeros((1, 7))
+ )
+ assert reads == [] and commits == []
+
+
+def test_root_read_after_only_joint_staging_uses_current_root_and_keeps_it_unwritten():
+ transaction, reads, commits = fixture()
+ ids = np.array([4, 1])
+ original = transaction._backend.get_entity_state
+
+ def current(entity):
+ state = original(entity)
+ state["root_pose"][:, 0] = np.arange(5) + 10
+ return state
+
+ transaction._backend.get_entity_state = current
+ with transaction.scoped(ids):
+ transaction.write_entity_state(
+ "object",
+ ids,
+ term_name="joint",
+ joint_names=("a",),
+ joint_positions=np.array([[0.2], [0.3]]),
+ )
+ assert reads == []
+ pose = transaction.read_entity_root_pose("object", ids)
+ np.testing.assert_array_equal(pose[:, 0], [14, 11])
+ pose[:] = 999
+ assert reads == ["object"]
+ assert commits[0].patches[0].root_pose is None
+ assert commits[0].patches[0].root_velocity is None
+ np.testing.assert_allclose(commits[0].patches[0].joint_positions, [[0.3], [0.2]])
diff --git a/tests/base/test_entity_scene_consumer.py b/tests/base/test_entity_scene_consumer.py
new file mode 100644
index 000000000..5e9380aa6
--- /dev/null
+++ b/tests/base/test_entity_scene_consumer.py
@@ -0,0 +1,196 @@
+"""Typed scene authoring and public-only multi-entity manager transactions."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from pathlib import Path
+
+import numpy as np
+import pytest
+from omegaconf import OmegaConf
+from unisim.entities import EntityVariantBinding, SceneEntitySpec
+
+from tests.envs.test_multi_entity_consumer import build_fixture_cfg
+from unilab.base import backend_factory
+from unilab.base.base import EnvCfg
+from unilab.base.config_materialization import apply_cfg_overrides
+from unilab.base.scene import SceneCfg
+from unilab.envs.manager_based_rl_env import make_manager_based_rl_env
+
+
+def test_hydra_materializes_entity_sources_binding_and_immutable_assignment():
+ scene = SceneCfg()
+ apply_cfg_overrides(
+ scene,
+ OmegaConf.create(
+ {
+ "entity_assets": [
+ {
+ "name": "object",
+ "source": {"model_file": "a.xml"},
+ "kind": "rigid",
+ "initial_state": {"position": [1.0, 2.0, 3.0]},
+ }
+ ],
+ "entity_variant": {
+ "target_entity": "object",
+ "plan": {
+ "variants": [{"model_file": "a.xml"}, {"model_file": "b.xml"}],
+ "assignment": [1, 1, 0],
+ "layout": "same_layout",
+ },
+ },
+ "entities": {
+ "tool": {"root_body_name": "object/base", "physical_entity": "object"}
+ },
+ }
+ ),
+ )
+ assert isinstance(scene.entity_assets[0], SceneEntitySpec)
+ assert scene.entity_assets[0].initial_state.position == (1.0, 2.0, 3.0)
+ assert isinstance(scene.entity_variant, EntityVariantBinding)
+ np.testing.assert_array_equal(scene.entity_variant.plan.assignment, [1, 1, 0])
+ assert not scene.entity_variant.plan.assignment.flags.writeable
+ assert scene.entities["tool"].physical_entity == "object"
+ with pytest.raises(ValueError):
+ SceneCfg(model_file="legacy.xml", entity_assets=scene.entity_assets)
+
+
+def test_asset_owner_sees_each_physical_and_variant_source_once(monkeypatch):
+ cfg = build_fixture_cfg(passive=False)
+ calls = []
+ monkeypatch.setattr(
+ backend_factory, "ensure_robot_assets_for_paths", lambda paths: calls.append(paths)
+ )
+ sentinel = object()
+ monkeypatch.setattr(backend_factory.unisim, "create_backend", lambda *args, **kwargs: sentinel)
+ assert backend_factory.create_backend("mujoco", cfg.scene, 2, 0.001) is sentinel
+ assert len(calls) == 1
+ expected = {entity.source.model_file for entity in cfg.scene.entity_assets if entity.source}
+ expected.update(source.model_file for source in cfg.scene.entity_variant.plan.variants)
+ assert expected.issubset(calls[0])
+ assert len(calls[0]) == len(set(calls[0]))
+
+
+@pytest.mark.parametrize("logical", ["robot", "object"])
+def test_logical_root_cannot_point_to_a_physical_entity_descendant(logical):
+ cfg = build_fixture_cfg(passive=True)
+ child = "finger" if logical == "robot" else "lid"
+ cfg.scene.entities[logical] = replace(
+ cfg.scene.entities[logical], root_body_name=f"{logical}/{child}"
+ )
+ with pytest.raises(ValueError, match=f"physical root '{logical}/base'"):
+ make_manager_based_rl_env(cfg, 2, "mujoco")
+
+
+def test_hydra_constructs_new_scene_without_early_unisim_tuple_validation():
+ cfg = EnvCfg()
+ apply_cfg_overrides(
+ cfg,
+ OmegaConf.create(
+ {
+ "scene": {
+ "_target_": "unilab.base.scene.SceneCfg",
+ "entity_assets": [
+ {
+ "name": "object",
+ "kind": "rigid",
+ "source": {"model_file": "object.xml"},
+ "initial_state": {"quaternion": [1.0, 0.0, 0.0, 0.0]},
+ }
+ ],
+ }
+ }
+ ),
+ )
+ assert isinstance(cfg.scene.entity_assets[0], SceneEntitySpec)
+ assert cfg.scene.entity_assets[0].initial_state.quaternion == (1.0, 0.0, 0.0, 0.0)
+
+
+def test_two_manager_entity_writes_commit_once_and_late_failure_commits_nothing(monkeypatch):
+ env = make_manager_based_rl_env(build_fixture_cfg(passive=False), 2, "mujoco")
+ try:
+ env.init_state()
+ commits = []
+ native = env._backend.reset_entities
+ monkeypatch.setattr(
+ env._backend,
+ "reset_entities",
+ lambda request: (commits.append(request), native(request))[-1],
+ )
+ ids = np.array([1])
+ with env._reset_state.scoped(ids):
+ env.scene["object"].write_root_link_pose_to_sim(
+ np.array([[0, 0, 2, 1, 0, 0, 0.0]]), env_ids=ids
+ )
+ env.scene["target"].write_root_link_pose_to_sim(
+ np.array([[3, 2, 1, 1, 0, 0, 0.0]]), env_ids=ids
+ )
+ assert len(commits) == 1
+ assert {patch.entity for patch in commits[0].patches} == {"object", "target"}
+ before = env._backend.get_state()
+ with pytest.raises(ValueError), env._reset_state.scoped(ids):
+ env.scene["object"].write_root_link_pose_to_sim(
+ np.array([[0, 0, 3, 1, 0, 0, 0.0]]), env_ids=ids
+ )
+ env.scene["target"].write_root_link_pose_to_sim(np.zeros((1, 7)), env_ids=ids)
+ assert len(commits) == 1
+ for key, value in before.items():
+ np.testing.assert_array_equal(env._backend.get_state()[key], value)
+ finally:
+ env.close()
+
+
+def test_different_entity_row_selections_fail_before_committing(monkeypatch):
+ env = make_manager_based_rl_env(build_fixture_cfg(passive=False), 2, "mujoco")
+ try:
+ commits = []
+ monkeypatch.setattr(env._backend, "reset_entities", lambda request: commits.append(request))
+ with pytest.raises(NotImplementedError, match="same selected env rows"):
+ with env._reset_state.scoped(np.array([0, 1])):
+ env.scene["object"].write_root_link_pose_to_sim(
+ np.array([[0, 0, 2, 1, 0, 0, 0.0]]), env_ids=np.array([1])
+ )
+ env.scene["target"].write_root_link_pose_to_sim(
+ np.array([[3, 2, 1, 1, 0, 0, 0.0]]), env_ids=np.array([0])
+ )
+ assert commits == []
+ finally:
+ env.close()
+
+
+def test_per_environment_entity_defaults_are_not_broadcast_from_first_variant():
+ cfg = build_fixture_cfg(passive=True)
+ variants = cfg.scene.entity_variant.plan.variants
+ for index, source in enumerate(variants):
+ path = Path(source.model_file)
+ text = path.read_text().replace(
+ "",
+ f'',
+ )
+ path.write_text(text)
+ cfg.scene.default_keyframe_name = "home"
+ # Every physical source needs the named key when selected by the composition owner.
+ for entity in cfg.scene.entity_assets:
+ if entity.name == "object":
+ continue
+ path = Path(entity.source.model_file)
+ key = (
+ ''
+ if entity.name == "robot"
+ else ''
+ )
+ text = path.read_text().replace("", f"{key}")
+ path.write_text(text)
+ env = make_manager_based_rl_env(cfg, 2, "mujoco")
+ try:
+ env.init_state()
+ np.testing.assert_allclose(env.scene["object"].data.default_joint_pos[:, 0], [0.1, 0.3])
+ np.testing.assert_allclose(env.scene["object"].data.joint_pos[:, 0], [0.1, 0.3])
+ np.testing.assert_allclose(env._control[:, 0], [0.37, 0.37])
+ env.step(np.full((2, 1), 0.1, dtype=np.float32))
+ env.reset(env_ids=np.array([1]))
+ np.testing.assert_allclose(env.scene["object"].data.joint_pos[1, 0], 0.3)
+ np.testing.assert_allclose(env._control[:, 0], [0.1, 0.37])
+ finally:
+ env.close()
diff --git a/tests/base/test_isaacsim_backend.py b/tests/base/test_isaacsim_backend.py
index 40c4421af..283ceb4da 100644
--- a/tests/base/test_isaacsim_backend.py
+++ b/tests/base/test_isaacsim_backend.py
@@ -31,10 +31,8 @@
build_worker_env,
resolve_isaacsim_runtime,
)
-from unisim.backend.isaacsim.worker import (
- _quat_rotate_wxyz,
- _resolve_articulation_root_prim_path,
-)
+from unisim.backend.isaacsim.scene_worker import _rotate
+from unisim.backend.isaacsim.worker import _resolve_articulation_root_prim_path
from unilab.base.backend_factory import create_backend
from unilab.base.base import EnvCfg
@@ -379,7 +377,7 @@ def test_interactive_renderer_roundtrip(scene_file: str, monkeypatch: pytest.Mon
backend.close()
-def test_interactive_playback_routes_startup_dimensions_and_camera(
+def test_interactive_playback_routes_dimensions_and_rejects_ignored_camera_options(
scene_file: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("DISPLAY", ":0")
@@ -406,11 +404,7 @@ def record_init_renderer(*args: Any, **kwargs: Any) -> None:
num_steps=1,
headless=False,
record_video=False,
- camera_kwargs={
- "cam_distance": 3.0,
- "cam_elevation": -15.0,
- "cam_azimuth": 45.0,
- },
+ camera_kwargs=CameraCfg(),
)
assert result is None
assert init_calls == [
@@ -418,13 +412,19 @@ def record_init_renderer(*args: Any, **kwargs: Any) -> None:
"headless": False,
"width": 64,
"height": 48,
- "camera_kwargs": CameraCfg(
- cam_distance=3.0,
- cam_elevation=-15.0,
- cam_azimuth=45.0,
- ),
+ "camera_kwargs": CameraCfg(),
}
]
+ with pytest.raises(NotImplementedError, match="interactive viewers"):
+ backend.run_playback(
+ env=SimpleNamespace(cfg=None),
+ initialize=lambda: 0,
+ step=lambda obs: obs + 1,
+ num_steps=1,
+ headless=False,
+ record_video=False,
+ camera_kwargs={"cam_distance": 3.0},
+ )
finally:
backend.close()
@@ -536,7 +536,7 @@ def test_root_angular_velocity_helper_converts_body_to_world() -> None:
half = np.sqrt(0.5)
quat = np.array([[half, half, 0.0, 0.0]], dtype=np.float32)
body_angvel = np.array([[0.0, -1.0, 0.0]], dtype=np.float32)
- np.testing.assert_allclose(_quat_rotate_wxyz(quat, body_angvel), [[0.0, 0.0, -1.0]], atol=1e-6)
+ np.testing.assert_allclose(_rotate(quat, body_angvel), [[0.0, 0.0, -1.0]], atol=1e-6)
class _FakePrim:
diff --git a/tests/envs/test_multi_entity_consumer.py b/tests/envs/test_multi_entity_consumer.py
new file mode 100644
index 000000000..d3684ad85
--- /dev/null
+++ b/tests/envs/test_multi_entity_consumer.py
@@ -0,0 +1,290 @@
+"""Same registered Manager-Based task consumes public UniSim entity contracts."""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import pickle
+import subprocess
+import sys
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import pytest
+from unisim.dr.types import FixedVariantPlan, ModelSourceDescriptor
+from unisim.entities import EntityInitialState, EntityVariantBinding, SceneEntitySpec
+
+from unilab.base import registry
+from unilab.base.entity import EntityCfg
+from unilab.base.env_factory import registry_env_factory
+from unilab.base.scene import SceneCfg
+from unilab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg, make_manager_based_rl_env
+from unilab.managers import (
+ ActionTerm,
+ ActionTermCfg,
+ EventTermCfg,
+ ObservationGroupCfg,
+ ObservationTermCfg,
+)
+
+TASK = "MultiEntityConsumerContract"
+_ASSETS = tempfile.TemporaryDirectory(prefix="unilab-m2-fixture-")
+
+
+def _sources(passive: bool):
+ root = Path(_ASSETS.name)
+ robot = root / "robot.xml"
+ robot.write_text("""
+
+
+
+ """)
+ variants = []
+ for index in range(2):
+ path = root / f"object-{passive}-{index}.xml"
+ child = (
+ (
+ ''
+ ''
+ )
+ if passive
+ else ""
+ )
+ path.write_text(f'''
+ {child}
+ ''')
+ variants.append(ModelSourceDescriptor(str(path)))
+ table = root / "table.xml"
+ table.write_text("""
+ """)
+ return robot, tuple(variants), table
+
+
+@dataclass(kw_only=True)
+class PositionActionCfg(ActionTermCfg):
+ def build(self, env):
+ return PositionAction(self, env)
+
+
+class PositionAction(ActionTerm):
+ def __init__(self, cfg, env):
+ super().__init__(cfg, env)
+ self.values = np.zeros((self.num_envs, 1), dtype=np.float32)
+
+ @property
+ def action_dim(self):
+ return 1
+
+ @property
+ def raw_action(self):
+ return self.values
+
+ def process_actions(self, actions):
+ self.values[:] = actions
+
+ def apply_actions(self):
+ self._entity.data.write_ctrl(self.values)
+
+
+def observations(env):
+ return np.concatenate(
+ (env.scene["robot"].data.joint_pos, env.scene["object"].data.root_link_pos_w), axis=1
+ )
+
+
+def reset_defaults(env, env_ids):
+ env.scene.reset_to_default(env_ids, term_name="fixture.defaults")
+
+
+def make_fixture_cfg():
+ return build_fixture_cfg(passive=True)
+
+
+def build_fixture_cfg(*, passive, num_envs=2):
+ robot, variants, table = _sources(passive)
+ entities = [
+ SceneEntitySpec(
+ "robot",
+ ModelSourceDescriptor(str(robot)),
+ root_mode="fixed",
+ initial_state=EntityInitialState(position=(-1.0, 0.0, 0.5)),
+ ),
+ SceneEntitySpec(
+ "object",
+ variants[0],
+ kind="articulation" if passive else "rigid",
+ initial_state=EntityInitialState(position=(0.0, 0.0, 1.0)),
+ ),
+ SceneEntitySpec(
+ "table", ModelSourceDescriptor(str(table)), kind="rigid", root_mode="fixed"
+ ),
+ ]
+ selectors = {
+ "robot": EntityCfg(
+ root_body_name="robot/base",
+ physical_entity="robot",
+ joint_names=("robot/hinge",),
+ actuator_names=("robot/drive",),
+ ),
+ "object": EntityCfg(
+ root_body_name="object/base",
+ physical_entity="object",
+ joint_names=("object/passive",) if passive else (),
+ actuator_names=(),
+ ),
+ "table": EntityCfg(
+ root_body_name="table/base", physical_entity="table", joint_names=(), actuator_names=()
+ ),
+ }
+ if not passive:
+ entities.append(
+ SceneEntitySpec(
+ "target",
+ kind="rigid",
+ root_mode="kinematic",
+ collision_enabled=False,
+ mirror_of="object",
+ initial_state=EntityInitialState(position=(2.0, 0.0, 1.0)),
+ )
+ )
+ selectors["target"] = EntityCfg(
+ root_body_name="target/base",
+ physical_entity="target",
+ joint_names=(),
+ actuator_names=(),
+ )
+ scene = SceneCfg(
+ entity_assets=tuple(entities),
+ entities=selectors,
+ primary_entity="robot",
+ entity_variant=EntityVariantBinding(
+ "object",
+ FixedVariantPlan(np.array([0, 1] if num_envs == 2 else [1, 1, 0, 1, 0]), variants),
+ ),
+ )
+ return ManagerBasedRlEnvCfg(
+ scene=scene,
+ sim_dt=0.001,
+ ctrl_dt=0.001,
+ max_episode_seconds=1.0,
+ observations={
+ "policy": ObservationGroupCfg(terms={"state": ObservationTermCfg(func=observations)})
+ },
+ actions={"position": PositionActionCfg(entity_name="robot")},
+ events={"defaults": EventTermCfg(func=reset_defaults, mode="reset")},
+ seed=1,
+ )
+
+
+registry.register_env_config(TASK, make_fixture_cfg)
+for _backend in ("mujoco", "isaacsim"):
+ registry.register_env(TASK, make_manager_based_rl_env, sim_backend=_backend)
+
+
+@pytest.mark.parametrize("num_envs", [2, 5])
+def test_registry_factory_mujoco_multi_entity_reset_isolation(num_envs):
+ factory = pickle.loads(pickle.dumps(registry_env_factory(TASK, "mujoco")))
+ env = factory(
+ num_envs=num_envs,
+ env_cfg_override={"scene": build_fixture_cfg(passive=True, num_envs=num_envs).scene},
+ )
+ try:
+ state = env.init_state()
+ assert state.obs["obs"].shape == (num_envs, 4)
+ assert env.action_space.shape == (1,)
+ assert env.scene["object"].data.joint_pos.shape == (num_envs, 1)
+ env.step(np.full((num_envs, 1), 0.2, dtype=np.float32))
+ robot_before = env.scene["robot"].data.joint_pos.copy()
+ object_before = env.scene["object"].data.root_link_pose_w.copy()
+ with env._reset_state.scoped(np.array([1])):
+ env.scene["object"].write_root_link_pose_to_sim(
+ np.array([[0.3, 0.4, 2.0, 1.0, 0.0, 0.0, 0.0]]), env_ids=np.array([1])
+ )
+ np.testing.assert_array_equal(env.scene["robot"].data.joint_pos, robot_before)
+ np.testing.assert_array_equal(
+ env.scene["object"].data.root_link_pose_w[0], object_before[0]
+ )
+ np.testing.assert_allclose(env.scene["object"].data.root_link_pos_w[1], [0.3, 0.4, 2.0])
+ with pytest.raises(ValueError), env._reset_state.scoped(np.array([1])):
+ env.scene["object"].write_root_link_pose_to_sim(np.zeros((1, 7)), env_ids=np.array([1]))
+ np.testing.assert_array_equal(env.scene["robot"].data.joint_pos, robot_before)
+ finally:
+ env.close()
+
+
+def test_factory_can_be_unpickled_in_a_fresh_process_without_parent_registry():
+ encoded = base64.b64encode(pickle.dumps(registry_env_factory(TASK, "mujoco"))).decode()
+ script = """
+import base64, pickle, sys
+factory = pickle.loads(base64.b64decode(sys.argv[1]))
+env = factory(num_envs=2)
+try:
+ state = env.init_state()
+ assert state.obs["obs"].shape == (2, 4)
+ assert env.action_space.shape == (1,)
+finally:
+ env.close()
+print("FACTORY_OK")
+"""
+ process_env = dict(os.environ)
+ # This validates physics/registry reconstruction, without a renderer. An
+ # inherited OSMesa choice otherwise makes SDK import require libOSMesa.
+ process_env["MUJOCO_GL"] = "disable"
+ process_env["UNILAB_EXTRA_REGISTRY_PACKAGES"] = __name__
+ result = subprocess.run(
+ [sys.executable, "-c", script, encoded],
+ env=process_env,
+ capture_output=True,
+ text=True,
+ timeout=30,
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+ assert "FACTORY_OK" in result.stdout
+
+
+@pytest.mark.skipif(
+ os.environ.get("UNILAB_TEST_M2_ISAACSIM") != "1", reason="explicit native runtime opt-in"
+)
+@pytest.mark.parametrize("passive", [True, False])
+def test_native_isaacsim_same_manager_task(passive, tmp_path):
+ factory = pickle.loads(pickle.dumps(registry_env_factory(TASK, "isaacsim")))
+ env = factory(num_envs=2, env_cfg_override={"scene": build_fixture_cfg(passive=passive).scene})
+ try:
+ state = env.init_state()
+ assert state.obs["obs"].shape == (2, 4)
+ assert env.action_space.shape == (1,)
+ env.step(np.zeros((2, 1), dtype=np.float32))
+ before = env.scene["robot"].data.joint_pos.copy()
+ with env._reset_state.scoped(np.array([1])):
+ env.scene["object"].write_root_link_pose_to_sim(
+ np.array([[0.3, 0.4, 2.0, 1.0, 0.0, 0.0, 0.0]]), env_ids=np.array([1])
+ )
+ np.testing.assert_array_equal(env.scene["robot"].data.joint_pos, before)
+ if not passive:
+ with env._reset_state.scoped(np.array([1])):
+ env.scene["target"].write_root_link_pose_to_sim(
+ np.array([[3.0, 2.0, 1.0, 1.0, 0.0, 0.0, 0.0]]), env_ids=np.array([1])
+ )
+ np.testing.assert_allclose(env.scene["target"].data.root_link_pos_w[1], [3, 2, 1])
+ (tmp_path / "evidence.json").write_text(
+ json.dumps(
+ {
+ "task": TASK,
+ "backend": "isaacsim",
+ "num_envs": 2,
+ "assignment": [0, 1],
+ "passive_articulation": passive,
+ "kinematic_mirror": not passive,
+ "obs_groups_spec": env.obs_groups_spec,
+ "action_shape": list(env.action_space.shape),
+ "object_root_positions": env.scene["object"].data.root_link_pos_w.tolist(),
+ "robot_joint_positions": env.scene["robot"].data.joint_pos.tolist(),
+ },
+ indent=2,
+ )
+ )
+ finally:
+ env.close()
diff --git a/uv.lock b/uv.lock
index d6919a5a9..db10bbfa6 100644
--- a/uv.lock
+++ b/uv.lock
@@ -5240,8 +5240,8 @@ requires-dist = [
{ name = "trimesh", marker = "extra == 'viser'", specifier = ">=3.21.7" },
{ name = "typing-extensions" },
{ name = "unilab-rl", specifier = "==1.2.1" },
- { name = "unisim-core", specifier = ">=1.4.2" },
- { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = ">=1.4.2" },
+ { name = "unisim-core", specifier = ">=1.5.0" },
+ { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = ">=1.5.0" },
{ name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" },
{ name = "wandb" },
{ name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" },
@@ -5284,13 +5284,13 @@ wheels = [
[[package]]
name = "unisim-core"
-version = "1.4.2"
+version = "1.5.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
{ name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/b9/e7/5954d490d66facf3449f1032d5cc13288fb4fd1166eaba4eea8987c4f4ca/unisim_core-1.4.2.tar.gz", hash = "sha256:e27d89292083cdff646d2dbeb0ac787ee4351e7ef811c20c6fed2a40578785b7", size = 269413, upload-time = "2026-09-16T04:22:39.708Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/1c/d8/e142031c7aeb29653ba04b882f5f8df57a7fc5603a7bae682b08121e005f/unisim_core-1.5.0.tar.gz", hash = "sha256:4fd0f276af7f1905812f9189e11422889993423e00a331abdecaa00a8cdd118c", size = 354365, upload-time = "2026-09-17T04:39:54.318Z" }
[package.optional-dependencies]
superdex = [
diff --git a/uv.rocm.lock b/uv.rocm.lock
index 7725f2723..5c9753d59 100644
--- a/uv.rocm.lock
+++ b/uv.rocm.lock
@@ -3790,7 +3790,7 @@ requires-dist = [
{ name = "triton-rocm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==3.6.0", index = "https://download.pytorch.org/whl/rocm7.2" },
{ name = "typing-extensions" },
{ name = "unilab-rl", specifier = "==1.2.1" },
- { name = "unisim-core", specifier = ">=1.4.2" },
+ { name = "unisim-core", specifier = ">=1.5.0" },
{ name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" },
{ name = "wandb" },
]
@@ -3830,13 +3830,13 @@ wheels = [
[[package]]
name = "unisim-core"
-version = "1.4.2"
+version = "1.5.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
{ name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/b9/e7/5954d490d66facf3449f1032d5cc13288fb4fd1166eaba4eea8987c4f4ca/unisim_core-1.4.2.tar.gz", hash = "sha256:e27d89292083cdff646d2dbeb0ac787ee4351e7ef811c20c6fed2a40578785b7", size = 269413, upload-time = "2026-09-16T04:22:39.708Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/1c/d8/e142031c7aeb29653ba04b882f5f8df57a7fc5603a7bae682b08121e005f/unisim_core-1.5.0.tar.gz", hash = "sha256:4fd0f276af7f1905812f9189e11422889993423e00a331abdecaa00a8cdd118c", size = 354365, upload-time = "2026-09-17T04:39:54.318Z" }
[[package]]
name = "urllib3"