"""Host-side IsaacSim/IsaacLab subprocess backend.
IsaacSim is deliberately kept out of the UniLab interpreter: the supported
IsaacSim 5.1 wheels require Python 3.11, while the main project supports a
different Python range. The NumPy-facing contract and lifecycle come from the
backend-neutral ``subprocess_ipc`` owner; this module supplies IsaacSim runtime
discovery, the worker entrypoint, clone-origin validation, and the eval-owned
Kit/viewer/camera capability boundary.
"""
from __future__ import annotations
import os
import time
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
from unisim.backend.base import (
CameraCfg,
PhysicsStateLayout,
PlayCapabilities,
PlayRenderPlan,
PublicStateWidths,
SelectedResetPublication,
SensorDescriptor,
TensorDataPlane,
TensorExecution,
TensorLifecycleCapabilities,
TensorProcessTopology,
TrackedBodyStateViews,
normalize_play_render_mode,
validate_torch_device,
)
from unisim.backend.isaacgym.backend import IsaacGymWorkerError
from unisim.backend.playback import display_available
from unisim.backend.subprocess_ipc import protocol
from unisim.backend.subprocess_ipc.backend import (
MjcfSubprocessBackend,
SubprocessModelInfo,
)
from unisim.backend.subprocess_ipc.isaac_common import (
mapped_default_actuator_rows,
mapped_default_body_rows,
mapped_default_dof_rows,
)
from unisim.backend.subprocess_ipc.scene_materialization import PreparedWorkerScene
from unisim.backend.subprocess_ipc.sensors import (
KIND_CONTACT_FORCE,
KIND_CONTACT_FOUND,
SceneSensorSpec,
UnsupportedSensorSpec,
)
from unisim.dr._host import interval_payload_to_host, reset_payload_to_host_dict
from unisim.dr.interval import (
INTERVAL_TERM_BODY_FORCE,
INTERVAL_TERM_BODY_TORQUE,
IntervalTermOp,
)
from unisim.dr.types import (
RESET_TERM_BASE_COM,
RESET_TERM_BASE_MASS,
RESET_TERM_BODY_INERTIA,
RESET_TERM_BODY_IPOS,
RESET_TERM_BODY_MASS,
RESET_TERM_DOF_ARMATURE,
RESET_TERM_DOF_DAMPING,
RESET_TERM_DOF_FRICTIONLOSS,
RESET_TERM_GEOM_FRICTION,
RESET_TERM_KD,
RESET_TERM_KP,
DomainRandomizationCapabilities,
FixedVariantLayout,
IntervalRandomizationPlan,
TensorResetRandomizationPayload,
_validate_reset_term,
validate_tensor_reset_payload,
)
from unisim.entities import SceneResetRequest
from unisim.progress import progress_enabled
from unisim.utils.quaternion import torch_quat_mul, torch_quat_rotate
from .dependencies import build_worker_env, resolve_isaacsim_runtime
from .physx_solver import PHYSX_SOLVER_AUTHORED_FIELDS, PhysxSolverConfig, solver_value_matches
from .raw_usd_cache import resolve_raw_usd_cache_root, resolve_role_usd_cache_root
from .tensor_ipc import (
ISAACSIM_CUDA_ATTACH,
ISAACSIM_CUDA_READY,
ISAACSIM_CUDA_RESET,
ISAACSIM_CUDA_STEP,
HostCudaIpcArena,
import_torch,
require_cuda_tensor,
require_cuda_tensor_of_dtype,
)
_MODULE_DIR = Path(__file__).resolve().parent
_WORKER_PATH = _MODULE_DIR / "worker.py"
_NATIVE_GEOMETRY_AUDIT_FIELDS = frozenset(
(
"schema_version",
"path",
"usd_authored_type_name",
"usd_schema_type_name",
"shape_kind",
"collision_enabled",
"physx_collision_api_present",
"mesh_approximation",
"contact_offset_authored",
"contact_offset_schema_value",
"contact_offset_simulation_determined",
"rest_offset_authored",
"rest_offset_schema_value",
"rest_offset_simulation_determined",
"local_transform",
"world_transform",
"local_bounds_empty",
"local_bounds_min",
"local_bounds_max",
"local_extent",
"world_bounds_empty",
"world_bounds_min",
"world_bounds_max",
"world_extent",
"extent_tolerance",
"zero_extent",
"zero_area",
"zero_volume_extent",
"planar_extent",
"mesh_vertex_count",
"mesh_unique_vertex_count",
"mesh_face_count",
"mesh_face_vertex_index_count",
"mesh_tolerance",
"mesh_all_points_coincident",
"mesh_points_min",
"mesh_points_max",
"mesh_points_extent",
"mesh_points_zero_extent",
"mesh_surface_area",
"mesh_zero_area",
)
)
[docs]
class IsaacSimRenderError(RuntimeError):
"""Raised when the requested IsaacSim render mode cannot be initialized."""
[docs]
class IsaacSimWorkerError(IsaacGymWorkerError):
"""Raised when the external IsaacSim/IsaacLab worker fails.
``IsaacGymWorkerError`` is retained as a compatibility ancestor because
early subprocess adapters exposed that exception as the public worker
failure type. The concrete class remains distinct, so callers can still
distinguish IsaacSim failures with an exact type check while existing
``except IsaacGymWorkerError`` handlers continue to work.
"""
[docs]
@dataclass(frozen=True)
class IsaacSimModelInfo(SubprocessModelInfo):
"""Opaque metadata returned by the IsaacSim worker handshake."""
[docs]
class IsaacSimBackend(MjcfSubprocessBackend):
"""Thin host client for the IsaacLab/PhysX worker.
The shared client owns pipe framing, shared-memory slot allocation,
timeout/crash diagnostics, XML cold-path metadata, and all NumPy state
views. Physics remains the default (``render_mode=None``). Eval/play
profiles pass a concrete render intent through the cold ``INIT`` payload
so Kit selects the correct experience before simulation materialization.
The worker owns all camera/viewport operations; this class validates the
NumPy-facing frame contract.
"""
_BACKEND_TYPE = "isaacsim"
_BACKEND_LABEL = "isaacsim"
_WORKER_ERROR_CLS = IsaacSimWorkerError
_MODEL_INFO_CLS = IsaacSimModelInfo
_tensor_cuda_ipc_requested = False
_worker_materialization_report: dict[str, Any] | None = None
[docs]
def __init__(
self,
scene: Any,
num_envs: int,
sim_dt: float,
*,
render_mode: str | None = None,
render_width: int = 1280,
render_height: int = 720,
solver_position_iteration_count: int | None = None,
solver_velocity_iteration_count: int | None = None,
bounce_threshold_velocity: float | None = None,
contact_offset: float | None = None,
rest_offset: float | None = None,
max_depenetration_velocity: float | None = None,
gpu_max_rigid_contact_count: int | None = None,
gpu_max_rigid_patch_count: int | None = None,
tensor_cuda_ipc: bool = False,
share_friction_materials: bool = False,
env_spacing: float | None = None,
**kwargs: Any,
) -> None:
mode = None if render_mode is None else normalize_play_render_mode(render_mode)
if isinstance(tensor_cuda_ipc, bool):
self._tensor_cuda_ipc_requested = tensor_cuda_ipc
else:
raise TypeError(f"tensor_cuda_ipc must be a boolean, got {tensor_cuda_ipc!r}")
if not isinstance(share_friction_materials, bool):
raise TypeError(
f"share_friction_materials must be a boolean, got {share_friction_materials!r}"
)
self._share_friction_materials = share_friction_materials
if env_spacing is not None:
if isinstance(env_spacing, bool) or not isinstance(env_spacing, (int, float)):
raise TypeError(f"env_spacing must be a positive number, got {env_spacing!r}")
if not np.isfinite(env_spacing) or float(env_spacing) <= 0.0:
raise ValueError(f"env_spacing must be finite and positive, got {env_spacing!r}")
self._env_spacing = float(env_spacing)
else:
self._env_spacing = 2.0
self._cuda_ipc_arena: HostCudaIpcArena | None = None
self._cuda_control_bounds: tuple[Any, Any] | None = None
self._cuda_reset_row_bounds: Any = None
self._cuda_reset_selected: Any = None
self._cuda_reset_true: Any = None
self._cuda_reset_sequence = 0
for name, value in (("render_width", render_width), ("render_height", render_height)):
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
raise ValueError(f"{name} must be a positive integer, got {value!r}")
# Validation is fail-closed at construction, before any worker spawn.
self._physx_solver = PhysxSolverConfig(
solver_position_iteration_count=solver_position_iteration_count,
solver_velocity_iteration_count=solver_velocity_iteration_count,
bounce_threshold_velocity=bounce_threshold_velocity,
contact_offset=contact_offset,
rest_offset=rest_offset,
max_depenetration_velocity=max_depenetration_velocity,
gpu_max_rigid_contact_count=gpu_max_rigid_contact_count,
gpu_max_rigid_patch_count=gpu_max_rigid_patch_count,
)
self._requested_render_mode = mode
self._resolved_render_mode: str | None = None
super().__init__(scene, num_envs, sim_dt, **kwargs)
if share_friction_materials and self._entity_scene is None:
raise NotImplementedError(
"IsaacSim friction material sharing requires a mapped Manager-Based scene"
)
if env_spacing is not None and self._entity_scene is None:
raise NotImplementedError(
"IsaacSim environment spacing requires a mapped Manager-Based scene"
)
self._render_width = int(render_width)
self._render_height = int(render_height)
self._staged_body_wrench = (
None
if self._entity_scene is None
else np.zeros((self._num_envs, self._entity_scene.layout.nbody, 6), dtype=np.float32)
)
self._body_wrench_pending = False
# The worker protocol carries no simulation clock, so playback snapshot
# time comes from this per-env host accumulator (advanced on public
# steps, zeroed on full state resets).
self._time_view = np.zeros((self._num_envs,), dtype=np.float64)
if self._tensor_cuda_ipc_requested and self._entity_scene is None:
raise NotImplementedError(
"IsaacSim CUDA IPC tensor lifecycle requires a mapped Manager-Based scene"
)
if self._entity_scene is not None:
self._validate_playback_joint_order()
self._model_file_playback_layout: PhysicsStateLayout | None = None
else:
self._model_file_playback_layout = self._resolve_model_file_playback_layout()
def _resolve_render_mode(self) -> str:
"""Resolve eval intent before Kit is launched."""
requested = self._requested_render_mode
if requested is None or requested == "none":
return "none"
if requested == "auto":
return "interactive" if display_available() else "record"
if requested == "interactive" and not display_available():
raise IsaacSimRenderError(
"IsaacSim interactive rendering was requested but no local display was found; "
"set DISPLAY/WAYLAND_DISPLAY or use training.play_render_mode=record for "
"headless RGB capture."
)
return requested
def _worker_init_payload(self) -> dict[str, Any]:
mode = self._resolve_render_mode()
self._resolved_render_mode = mode
raw_usd_cache_root = resolve_raw_usd_cache_root()
role_usd_cache_root = resolve_role_usd_cache_root()
return {
"render_mode": mode,
"render_width": self._render_width,
"render_height": self._render_height,
# Progress frames interleave with the INIT reply; only opt in when
# the host will render them so raw-protocol clients never see them.
"init_progress": progress_enabled(),
"contact_force_sensors": self._contact_force_sensor_payload(),
"body_net_contact_entities": self._body_net_contact_entity_payload(),
"physx_solver": self._physx_solver.to_payload(),
"raw_usd_cache_dir": (None if raw_usd_cache_root is None else str(raw_usd_cache_root)),
"role_usd_cache_dir": (
None if role_usd_cache_root is None else str(role_usd_cache_root)
),
"tensor_cuda_ipc": getattr(self, "_tensor_cuda_ipc_requested", False),
"share_friction_materials": getattr(self, "_share_friction_materials", False),
"env_spacing": getattr(self, "_env_spacing", 2.0),
}
def _worker_configuration_requested(self) -> dict[str, Any]:
return self._physx_solver.to_payload()
def _allocate_slots(self) -> None:
"""Keep the opt-in tensor lifecycle off the model-file CPU shm plane."""
if not self._tensor_cuda_ipc_requested:
super()._allocate_slots()
return
# The shared initializer still seeds initial control metadata through
# ``_slots['ctrl']``. Use a bootstrap-only array, never a shared-memory
# segment, and clear it in ``_capture_entity_report`` at the cold-path
# boundary so model-file state access fails closed instead of reading stale
# host storage.
self._shm_handles.clear()
self._slots.clear()
self._slots["ctrl"] = np.zeros((self._num_envs, self.num_actuators), dtype=np.float32)
def _slot_specs(self) -> dict[str, dict[str, Any]]:
if self._tensor_cuda_ipc_requested:
return {}
return super()._slot_specs()
def _capture_entity_report(self, meta: dict[str, Any]) -> None:
super()._capture_entity_report(meta)
if self._tensor_cuda_ipc_requested:
self._slots.clear()
def _require_state(self, operation: str) -> None:
"""Allow the CUDA IPC lifecycle after bootstrap CPU slots are cleared."""
self._require_materialized()
if self._worker_dead_error is not None:
raise self._worker_error(
f"isaacsim worker is unavailable from an earlier failure; refusing {operation}"
) from self._worker_dead_error
if not self._slots and not self._tensor_cuda_ipc_requested:
raise self._worker_error(
f"isaacsim backend is closed or not materialized; refusing {operation}"
)
def _resolve_worker_entity_gravity(self, prepared: PreparedWorkerScene) -> None:
"""Resolve unset gravity requests to the historical IsaacSim role default."""
for entry in prepared.payload["scene_entities"]:
if entry["gravity_disabled"] is None:
entry["gravity_disabled"] = (
entry["root_mode"] == "kinematic"
or entry["kind"] == "rigid"
and entry["root_mode"] == "fixed"
)
_MAPPED_SUPPORTED_RESET_TERMS = frozenset(
{
RESET_TERM_GEOM_FRICTION,
RESET_TERM_KP,
RESET_TERM_KD,
RESET_TERM_BODY_MASS,
RESET_TERM_BODY_INERTIA,
RESET_TERM_BODY_IPOS,
RESET_TERM_BASE_MASS,
RESET_TERM_BASE_COM,
RESET_TERM_DOF_DAMPING,
RESET_TERM_DOF_ARMATURE,
RESET_TERM_DOF_FRICTIONLOSS,
}
)
_FIXED_VARIANT_LAYOUTS = frozenset(
{
FixedVariantLayout.SAME_LAYOUT,
FixedVariantLayout.UNIFORM_PUBLIC_LAYOUT,
}
)
[docs]
def get_dr_capabilities(self) -> DomainRandomizationCapabilities:
"""Advertise only reset and wrench terms implemented by mapped scenes.
Host-plane sets cover ``torch.cpu`` payloads; the device plane
(``torch.cuda``) carries the CUDA IPC PD-gain lane only.
"""
if self._entity_scene is None:
return super().get_dr_capabilities()
terms = frozenset({INTERVAL_TERM_BODY_FORCE, INTERVAL_TERM_BODY_TORQUE})
has_variants = self._entity_scene.owner.variant_plan is not None
reset_terms = self._MAPPED_SUPPORTED_RESET_TERMS
if getattr(self, "_share_friction_materials", False):
reset_terms = reset_terms - {RESET_TERM_GEOM_FRICTION}
return DomainRandomizationCapabilities(
supported_interval_terms=terms,
supported_reset_terms=reset_terms,
supported_device_reset_terms=frozenset({RESET_TERM_KP, RESET_TERM_KD}),
supports_fixed_variants=has_variants,
supported_fixed_variant_layouts=(
self._FIXED_VARIANT_LAYOUTS if has_variants else frozenset()
),
supports_per_env_playback=has_variants,
)
@staticmethod
def _coerce_mapped_reset_field(values: Any, name: str, shape: tuple[int, ...]) -> np.ndarray:
if not isinstance(values, np.ndarray) or not np.issubdtype(values.dtype, np.floating):
raise TypeError(f"isaacsim {name} must be a floating NumPy array")
array = np.asarray(values, dtype=np.float32)
if array.shape != shape:
raise ValueError(f"isaacsim {name} must have shape {shape}, got {array.shape}")
if not np.isfinite(array).all() or np.any(
np.abs(array.astype(np.float64)) > np.finfo(np.float32).max
):
raise ValueError(f"isaacsim {name} must contain finite float32 values")
return array.copy()
_MAPPED_UNSUPPORTED_RESET_REASONS = {
"gravity": "PhysX scene gravity is simulation-global, not per-environment",
"body_iquat": "PhysX derives the body principal-axes frame from the inertia tensor",
"geom_size": "PhysX collision geometry is immutable after materialization",
"geom_solref": "PhysX exposes no per-geom solver time-constant equivalent",
"geom_solimp": "PhysX exposes no per-geom solver impedance equivalent",
RESET_TERM_GEOM_FRICTION: (
"shared friction materials make reset-time sliding friction immutable; "
"disable share_friction_materials for per-geom friction DR"
),
}
def _mapped_base_body_column(self) -> int:
scene = self._require_mapped_entity_scene()
entity = scene.layout.entities[self._primary_entity_index()]
return entity.body_ids[entity.body_names.index(entity.root_body)]
def _mapped_owned_body_columns(self) -> np.ndarray:
"""Public body columns owned by an entity (excludes world-body rows)."""
layout = self._require_mapped_entity_scene().layout
return np.sort(
np.concatenate(
[np.asarray(list(entity.body_ids), dtype=np.int64) for entity in layout.entities]
)
)
def _mapped_default_body_rows(
self, field: str, rows: np.ndarray, width: int | None = None
) -> np.ndarray:
"""Build per-row variant-assigned public body-property defaults."""
return mapped_default_body_rows(
self, self._canonical_body_table(field, width), field, rows, width
)
def _mapped_default_geom_friction(self, rows: np.ndarray) -> np.ndarray:
"""Build per-row variant-assigned public geom-friction defaults.
uniform_public_layout variants may omit optional slots, so each
variant's rows scatter into the padded public table through
order-preserving name matching; omitted slots keep the zero default.
"""
scene = self._require_mapped_entity_scene()
layout = scene.layout
values = np.zeros((rows.size, layout.ngeom, 3), dtype=np.float32)
offset = 0
for entity, entry in zip(layout.entities, scene.payload["scene_entities"]):
count = len(entity.geoms)
if count:
if all(len(record["geom_names"]) == count for record in entry["variants"]):
defaults = np.asarray(
[
entry["variants"][int(entry["assignment"][int(env)])]["geom_friction"]
for env in rows
],
dtype=np.float32,
)
values[:, offset : offset + count, 0] = defaults[:, :, 0]
values[:, offset : offset + count, 1] = defaults[:, :, 0]
else:
public_names = [geom.name for geom in entity.geoms]
for row, env in enumerate(rows):
record = entry["variants"][int(entry["assignment"][int(env)])]
names = list(record["geom_names"])
friction = np.asarray(record["geom_friction"], dtype=np.float32)
if (
friction.shape != (len(names), 3)
or not np.isfinite(friction).all()
or np.any(friction < 0.0)
):
raise self._worker_error(
f"compiled variant geom_friction is malformed for entity "
f"{entity.name}"
)
slots: list[int] = []
cursor = 0
for name in names:
while cursor < count and public_names[cursor] != name:
cursor += 1
if cursor == count:
raise self._worker_error(
f"compiled variant geom_friction is malformed for entity "
f"{entity.name}"
)
slots.append(cursor)
cursor += 1
columns = [offset + slot for slot in slots]
values[row, columns, 0] = friction[:, 0]
values[row, columns, 1] = friction[:, 0]
offset += count
return values
def _mapped_default_actuator_rows(self, field: str, rows: np.ndarray) -> np.ndarray:
"""Build per-row variant-assigned public actuator-gain defaults."""
return mapped_default_actuator_rows(self, field, rows)
def _mapped_default_dof_rows(self, field: str, rows: np.ndarray) -> np.ndarray:
"""Build per-row variant-assigned public DOF-property defaults."""
return mapped_default_dof_rows(self, field, rows)
def _check_mapped_body_columns(
self,
field: str,
values: np.ndarray,
defaults: np.ndarray,
width: int | None = None,
) -> None:
"""Fail closed on body columns the mapped worker cannot mutate."""
scene = self._require_mapped_entity_scene()
layout = scene.layout
canonical = self._canonical_body_table(field, width)
owned = np.zeros(layout.nbody, dtype=bool)
for entity in layout.entities:
columns = list(entity.body_ids)
owned[columns] = True
if entity.root_mode != "kinematic":
continue
actual = values[:, columns] if width is None else values[:, columns, :]
expected = defaults[:, columns] if width is None else defaults[:, columns, :]
if not np.allclose(actual, expected, rtol=1e-5, atol=1e-8):
raise ValueError(
f"isaacsim {field} cannot mutate kinematic entity {entity.name} columns; "
"visual mirrors keep their variant defaults"
)
if not owned.all():
missing = np.flatnonzero(~owned)
actual = values[:, missing] if width is None else values[:, missing, :]
expected = np.broadcast_to(
canonical[missing],
(values.shape[0], missing.size)
if width is None
else (values.shape[0], missing.size, width),
)
if not np.allclose(actual, expected, rtol=1e-5, atol=1e-8):
raise ValueError(
f"isaacsim {field} columns {missing.tolist()} belong to no entity and "
"must keep canonical defaults"
)
def _check_mapped_actuator_columns(
self, term: str, values: np.ndarray, defaults: np.ndarray
) -> None:
"""Fail closed on actuator columns of kinematic mirror entities."""
scene = self._require_mapped_entity_scene()
for entity in scene.layout.entities:
if entity.root_mode != "kinematic" or not entity.actuator_indices:
continue
columns = list(entity.actuator_indices)
if not np.allclose(values[:, columns], defaults[:, columns], rtol=1e-5, atol=1e-8):
raise ValueError(
f"isaacsim {term} cannot mutate kinematic entity {entity.name} actuators; "
"visual mirrors keep their variant defaults"
)
def _check_mapped_dof_columns(
self, term: str, values: np.ndarray, defaults: np.ndarray
) -> None:
"""Fail closed on DOF columns of kinematic mirror entities."""
scene = self._require_mapped_entity_scene()
for entity in scene.layout.entities:
if entity.root_mode != "kinematic" or not entity.joints:
continue
columns = [joint.qvel_indices[0] for joint in entity.joints]
if not np.allclose(values[:, columns], defaults[:, columns], rtol=1e-5, atol=1e-8):
raise ValueError(
f"isaacsim {term} cannot mutate kinematic entity {entity.name} joints; "
"visual mirrors keep their variant defaults"
)
def _validated_mapped_reset_randomization(
self, randomization: TensorResetRandomizationPayload | None, rows: np.ndarray
) -> TensorResetRandomizationPayload | None:
if randomization is None:
return None
if not isinstance(randomization, TensorResetRandomizationPayload):
raise TypeError(
"isaacsim reset randomization must be a TensorResetRandomizationPayload"
)
if randomization.is_empty():
return None
scene = self._require_mapped_entity_scene()
layout = scene.layout
# Host plane: the mapped subprocess worker is a NumPy boundary, so the
# torch-first payload must carry torch.cpu tensors.
validate_tensor_reset_payload(randomization, device="cpu")
unsupported = self.get_dr_capabilities().get_unsupported_reset_terms(
randomization.requested_terms(), plane="host"
)
if unsupported:
details = "; ".join(
f"{term}: {self._MAPPED_UNSUPPORTED_RESET_REASONS.get(term, term)}"
for term in sorted(unsupported)
)
raise NotImplementedError(
f"isaacsim does not support reset domain randomization terms: {details}."
)
count = rows.size
host_values = reset_payload_to_host_dict(randomization)
geom_friction: np.ndarray | None = None
if RESET_TERM_GEOM_FRICTION in host_values:
geom_friction = self._coerce_mapped_reset_field(
host_values[RESET_TERM_GEOM_FRICTION],
"geom_friction",
(count, layout.ngeom, 3),
)
if np.any(geom_friction < 0.0):
raise ValueError("isaacsim geom_friction values must be nonnegative")
if np.any(geom_friction[..., 0] != geom_friction[..., 1]) or np.any(
geom_friction[..., 2] != 0.0
):
raise ValueError(
"isaacsim geom_friction requires static == dynamic and a zero third column"
)
body_mass: np.ndarray | None = None
if (
RESET_TERM_BODY_MASS in host_values
or RESET_TERM_BASE_MASS in host_values
):
body_mass = (
self._coerce_mapped_reset_field(
host_values[RESET_TERM_BODY_MASS], "body_mass", (count, layout.nbody)
)
if RESET_TERM_BODY_MASS in host_values
else self._mapped_default_body_rows("body_mass", rows)
)
if RESET_TERM_BASE_MASS in host_values:
# Adapter-private host-side delta merge: final = baseline + delta.
delta = self._coerce_mapped_reset_field(
host_values[RESET_TERM_BASE_MASS], "base_mass_delta", (count,)
)
body_mass[:, self._mapped_base_body_column()] += delta
# Unowned public rows (for example the world body) carry zero
# canonical placeholders and are pinned to defaults separately.
if np.any(body_mass[:, self._mapped_owned_body_columns()] <= 0.0):
raise ValueError("isaacsim body_mass values must be strictly positive")
self._check_mapped_body_columns(
"body_mass", body_mass, self._mapped_default_body_rows("body_mass", rows)
)
body_ipos: np.ndarray | None = None
if RESET_TERM_BODY_IPOS in host_values or RESET_TERM_BASE_COM in host_values:
body_ipos = (
self._coerce_mapped_reset_field(
host_values[RESET_TERM_BODY_IPOS], "body_ipos", (count, layout.nbody, 3)
)
if RESET_TERM_BODY_IPOS in host_values
else self._mapped_default_body_rows("body_ipos", rows, width=3)
)
if RESET_TERM_BASE_COM in host_values:
# Adapter-private host-side delta merge: final = baseline + delta.
offset = self._coerce_mapped_reset_field(
host_values[RESET_TERM_BASE_COM], "base_com_offset", (count, 3)
)
body_ipos[:, self._mapped_base_body_column(), :] += offset
self._check_mapped_body_columns(
"body_ipos",
body_ipos,
self._mapped_default_body_rows("body_ipos", rows, width=3),
width=3,
)
body_inertia: np.ndarray | None = None
if RESET_TERM_BODY_INERTIA in host_values:
body_inertia = self._coerce_mapped_reset_field(
host_values[RESET_TERM_BODY_INERTIA], "body_inertia", (count, layout.nbody, 3)
)
if np.any(body_inertia[:, self._mapped_owned_body_columns(), :] <= 0.0):
raise ValueError("isaacsim body_inertia values must be strictly positive")
self._check_mapped_body_columns(
"body_inertia",
body_inertia,
self._mapped_default_body_rows("body_inertia", rows, width=3),
width=3,
)
kp: np.ndarray | None = None
if RESET_TERM_KP in host_values:
kp = self._coerce_mapped_reset_field(
host_values[RESET_TERM_KP], "kp", (count, layout.nu)
)
if np.any(kp < 0.0):
raise ValueError("isaacsim kp values must be nonnegative")
self._check_mapped_actuator_columns(
"kp", kp, self._mapped_default_actuator_rows("dof_stiffness", rows)
)
kd: np.ndarray | None = None
if RESET_TERM_KD in host_values:
kd = self._coerce_mapped_reset_field(
host_values[RESET_TERM_KD], "kd", (count, layout.nu)
)
if np.any(kd < 0.0):
raise ValueError("isaacsim kd values must be nonnegative")
self._check_mapped_actuator_columns(
"kd", kd, self._mapped_default_actuator_rows("dof_damping", rows)
)
root_columns = [column for entity in layout.entities for column in entity.root_qvel_indices]
dof_values: dict[str, np.ndarray | None] = {}
for term, record_field in (
("dof_damping", None),
("dof_armature", "dof_armature"),
("dof_frictionloss", "dof_friction"),
):
requested = host_values.get(term)
if requested is None:
dof_values[term] = None
continue
values = self._coerce_mapped_reset_field(requested, term, (count, layout.nv))
if np.any(values < 0.0):
raise ValueError(f"isaacsim {term} values must be nonnegative")
if root_columns and np.any(values[:, root_columns] != 0.0):
raise ValueError(
f"isaacsim {term} free-root columns must remain zero; PhysX exposes no "
"root DOF damping, armature or friction"
)
defaults = (
np.zeros((count, layout.nv), dtype=np.float32)
if record_field is None
else self._mapped_default_dof_rows(record_field, rows)
)
self._check_mapped_dof_columns(term, values, defaults)
dof_values[term] = values
import torch
# The delta terms are consumed by the merge above; the returned payload
# carries final absolute host-plane tables only.
return TensorResetRandomizationPayload(
**{
term: torch.from_numpy(values)
for term, values in (
(RESET_TERM_BODY_MASS, body_mass),
(RESET_TERM_BODY_IPOS, body_ipos),
(RESET_TERM_BODY_INERTIA, body_inertia),
(RESET_TERM_GEOM_FRICTION, geom_friction),
(RESET_TERM_KP, kp),
(RESET_TERM_KD, kd),
("dof_damping", dof_values["dof_damping"]),
("dof_armature", dof_values["dof_armature"]),
("dof_frictionloss", dof_values["dof_frictionloss"]),
)
if values is not None
}
)
[docs]
def get_reset_term_default(self, term: str) -> np.ndarray:
"""Return per-environment variant-assigned mapped reset defaults."""
if self._entity_scene is None:
return super().get_reset_term_default(term)
_validate_reset_term(term)
if not self.get_dr_capabilities().supports_reset_term(term):
raise NotImplementedError(f"IsaacSimBackend does not support reset term {term!r}")
rows = np.arange(self._num_envs, dtype=np.intp)
layout = self._entity_scene.layout
if term == RESET_TERM_BASE_MASS:
value: np.ndarray = np.zeros((self._num_envs,), dtype=np.float32)
elif term == RESET_TERM_BASE_COM:
value = np.zeros((self._num_envs, 3), dtype=np.float32)
elif term == RESET_TERM_BODY_MASS:
value = self._mapped_default_body_rows("body_mass", rows)
elif term == RESET_TERM_BODY_IPOS:
value = self._mapped_default_body_rows("body_ipos", rows, width=3)
elif term == RESET_TERM_BODY_INERTIA:
value = self._mapped_default_body_rows("body_inertia", rows, width=3)
elif term == RESET_TERM_GEOM_FRICTION:
value = self._mapped_default_geom_friction(rows)
elif term == RESET_TERM_KP:
value = self._mapped_default_actuator_rows("dof_stiffness", rows)
elif term == RESET_TERM_KD:
value = self._mapped_default_actuator_rows("dof_damping", rows)
elif term == RESET_TERM_DOF_DAMPING:
value = np.zeros((self._num_envs, layout.nv), dtype=np.float32)
elif term == RESET_TERM_DOF_ARMATURE:
value = self._mapped_default_dof_rows("dof_armature", rows)
elif term == RESET_TERM_DOF_FRICTIONLOSS:
value = self._mapped_default_dof_rows("dof_friction", rows)
else: # pragma: no cover - guarded by supports_reset_term
raise NotImplementedError(f"IsaacSimBackend does not support reset term {term!r}")
result = np.array(value, dtype=np.float32, copy=True)
result.setflags(write=False)
return result
[docs]
def apply_interval_randomization(self, plan: IntervalRandomizationPlan) -> None:
if plan.is_empty():
return
if self._entity_scene is not None:
self._reject_wrench_write_inside_pre_step_control("apply_interval_randomization")
assert self._staged_body_wrench is not None
self._staged_body_wrench.fill(0.0)
self._body_wrench_pending = False
super().apply_interval_randomization(plan)
def _interval_term_handlers(self) -> dict[str, Callable[[IntervalTermOp], None]]:
if self._entity_scene is None:
return super()._interval_term_handlers()
return {
INTERVAL_TERM_BODY_FORCE: self._apply_interval_body_force,
INTERVAL_TERM_BODY_TORQUE: self._apply_interval_body_torque,
}
def _apply_interval_body_force(self, op: IntervalTermOp) -> None:
if op.body_ids is None:
raise ValueError("interval body force requires body_ids")
self.apply_body_force(op.body_ids, interval_payload_to_host(op.payload))
def _apply_interval_body_torque(self, op: IntervalTermOp) -> None:
if op.body_ids is None:
raise ValueError("interval body torque requires body_ids")
self._apply_body_torque(op.body_ids, interval_payload_to_host(op.payload))
def _apply_body_torque(self, body_ids: np.ndarray, torque: np.ndarray) -> None:
zero_force = np.zeros((self._num_envs, len(body_ids), 3), dtype=np.float32)
self.apply_body_force(body_ids, zero_force, torque=torque)
[docs]
def apply_body_force(
self,
body_ids: np.ndarray,
force: np.ndarray,
torque: np.ndarray | None = None,
) -> None:
"""Stage a world-frame wrench for the next mapped-scene step.
The force acts at each target body's center of mass and the optional
torque is about that center. Repeated submissions accumulate until the
next successful step consumes them.
"""
if self._entity_scene is None:
super().apply_body_force(body_ids, force, torque)
return
self._reject_wrench_write_inside_pre_step_control("apply_body_force")
self._require_state("interval body force perturbation")
assert self._staged_body_wrench is not None
body_ids_np = np.asarray(body_ids)
if (
body_ids_np.ndim != 1
or body_ids_np.dtype.kind not in "iu"
or np.issubdtype(body_ids_np.dtype, np.bool_)
):
raise ValueError("body_ids must be a one-dimensional integer array")
if np.any(body_ids_np < 0) or np.any(body_ids_np >= self._staged_body_wrench.shape[1]):
raise ValueError(
f"body ids must be in [0, {self._staged_body_wrench.shape[1]}), "
f"got range [{body_ids_np.min(initial=0)}, {body_ids_np.max(initial=0)}]"
)
body_ids_np = body_ids_np.astype(np.intp, copy=False)
expected_shape = (self._num_envs, body_ids_np.size, 3)
force_np = np.asarray(force, dtype=np.float32)
if force_np.shape != expected_shape:
raise ValueError(f"body force must have shape {expected_shape}, got {force_np.shape}")
if not np.isfinite(force_np).all():
raise ValueError("body force contains NaN or Inf")
torque_np = None
if torque is not None:
torque_np = np.asarray(torque, dtype=np.float32)
if torque_np.shape != expected_shape:
raise ValueError(
f"body torque must have shape {expected_shape}, got {torque_np.shape}"
)
if not np.isfinite(torque_np).all():
raise ValueError("body torque contains NaN or Inf")
for body_offset, body_id in enumerate(body_ids_np):
self._staged_body_wrench[:, int(body_id), 0:3] += force_np[:, body_offset, :]
if torque_np is not None:
self._staged_body_wrench[:, int(body_id), 3:6] += torque_np[:, body_offset, :]
self._body_wrench_pending = bool(
np.any(self._staged_body_wrench) or self._body_wrench_pending
)
def _step_payload(
self,
nsteps: int,
*,
body_wrench: np.ndarray | None = None,
) -> dict[str, Any]:
payload = super()._step_payload(nsteps)
if body_wrench is not None:
payload["body_wrench"] = body_wrench.tobytes(order="C")
elif self._entity_scene is not None and self._body_wrench_pending:
assert self._staged_body_wrench is not None
# NumPy pickle internals are not stable across the host and Isaac
# worker interpreter versions; raw C-order bytes are.
payload["body_wrench"] = self._staged_body_wrench.tobytes(order="C")
return payload
[docs]
def set_pre_step_control(self, fn: Any | None) -> None:
"""Register a callback for mapped scenes; model-file callbacks stay rejected."""
if fn is not None and self._entity_scene is None:
super().set_pre_step_control(fn)
return
self._pre_step_control_fn = fn
def _host_step(self, ctrl: np.ndarray, nsteps: int = 1) -> dict[str, dict[str, float]]:
if self._entity_scene is None or self._pre_step_control_fn is None:
result = super()._host_step(ctrl, nsteps)
else:
self._require_state("step")
if isinstance(nsteps, bool) or int(nsteps) <= 0:
raise ValueError(f"nsteps must be a positive integer, got {nsteps!r}")
self._require_materialized()
ctrl_array = np.asarray(ctrl, dtype=np.float32)
expected = (self._num_envs, self.num_actuators)
if ctrl_array.shape != expected:
raise ValueError(f"ctrl must have shape {expected}, got {ctrl_array.shape}")
if not np.isfinite(ctrl_array).all():
raise ValueError("control must contain finite target values")
ctrl_array = np.clip(
ctrl_array, self._entity_scene.control_lower, self._entity_scene.control_upper
)
result = self._step_with_pre_step_control(ctrl_array, int(nsteps))
self._advance_playback_time(int(nsteps))
return result
def _advance_playback_time(self, nsteps: int) -> None:
"""Advance the host-side playback clock after a completed public step."""
self._playback_time_view()[:] += nsteps * self._sim_dt
def _playback_time_view(self) -> np.ndarray:
"""Return the constructor-owned per-env playback clock."""
return self._time_view
def _step_with_pre_step_control(
self, ctrl: np.ndarray, nsteps: int
) -> dict[str, dict[str, float]]:
"""Run one mapped worker STEP per public physics substep.
The worker owns each native substep, while this host loop preserves the
backend-neutral callback boundary: shared state is refreshed after every
one-step command before the next callback invocation. This is a
correctness path, not the device-resident controller investigated by
#152.
"""
assert self._entity_scene is not None
assert self._staged_body_wrench is not None
fixed_wrench = self._staged_body_wrench.copy()
composed_wrench = np.zeros_like(fixed_wrench)
control_ms = 0.0
physics_ms = 0.0
started = time.perf_counter()
self._pre_step_control_active = True
try:
for substep in range(nsteps):
callback_started = time.perf_counter()
output = self._convert_pre_step_control(ctrl)
converted_ctrl = np.clip(
output.ctrl,
self._entity_scene.control_lower,
self._entity_scene.control_upper,
)
if not np.isfinite(converted_ctrl).all():
raise ValueError("pre-step control must contain finite target values")
composed_wrench[:] = fixed_wrench
if output.force is not None or output.torque is not None:
body_ids = np.asarray(output.body_ids, dtype=np.intp)
if np.any(body_ids < 0) or np.any(body_ids >= composed_wrench.shape[1]):
raise ValueError(
f"pre-step control body ids must be in "
f"[0, {composed_wrench.shape[1]}), got range "
f"[{body_ids.min(initial=0)}, {body_ids.max(initial=0)}]"
)
for body_offset, body_id in enumerate(body_ids):
if output.force is not None:
composed_wrench[:, int(body_id), 0:3] += output.force[
:, body_offset, :
].astype(np.float32, copy=False)
if output.torque is not None:
composed_wrench[:, int(body_id), 3:6] += output.torque[
:, body_offset, :
].astype(np.float32, copy=False)
np.copyto(self._slots["ctrl"], converted_ctrl)
control_ms += time.perf_counter() - callback_started
step_payload = self._step_payload(1, body_wrench=composed_wrench)
payload = self._request(protocol.CMD_STEP, step_payload, expect=protocol.CMD_READY)
if isinstance(payload, dict):
physics_ms += float(payload.get("timing", {}).get("physics_ms", 0.0))
if substep == nsteps - 1:
self._after_step(step_payload)
self._stale_body_ids.clear()
finally:
self._pre_step_control_active = False
# A callback step consumes staged interval wrench even if a later
# callback or worker substep fails; completed native substeps cannot
# be rolled back and a partially applied wrench must not leak.
self._staged_body_wrench.fill(0.0)
self._body_wrench_pending = False
return {
"timing": {
"control_upload_ms": control_ms * 1000.0,
"physics_ms": physics_ms,
"worker_ipc_total_ms": (time.perf_counter() - started) * 1000.0,
}
}
def _after_step(self, payload: dict[str, Any]) -> None:
if "body_wrench" in payload:
assert self._staged_body_wrench is not None
self._staged_body_wrench.fill(0.0)
self._body_wrench_pending = False
def _host_reset(self, env_ids: np.ndarray | None = None) -> None:
if self._tensor_cuda_ipc_requested and self._entity_scene is not None:
if env_ids is not None:
raise NotImplementedError(
"IsaacSim CUDA IPC model-file selected reset is unsupported; "
"use set_state_tensor"
)
prepared = self._entity_scene
torch = import_torch()
arena = self._ensure_cuda_ipc_arena()
device = arena.qpos.device
rows = torch.arange(self._num_envs, dtype=torch.int64, device=device)
qpos = torch.as_tensor(
np.asarray(prepared.qpos, dtype=np.float32),
dtype=torch.float32,
device=device,
)
qvel = torch.as_tensor(
np.asarray(prepared.qvel, dtype=np.float32),
dtype=torch.float32,
device=device,
)
self.set_state_tensor(rows, qpos, qvel)
return
super()._host_reset(env_ids)
if self._entity_scene is None:
return
rows = np.arange(self._num_envs) if env_ids is None else np.asarray(env_ids)
if rows.ndim != 1 or rows.dtype.kind not in "iu":
raise ValueError("reset env_ids must be one-dimensional integers")
if not len(rows):
return
if (
np.any(rows < 0)
or np.any(rows >= self._num_envs)
or len(set(rows.tolist())) != len(rows)
):
raise ValueError("reset env_ids must be distinct and in range")
assert self._staged_body_wrench is not None
# A full default-state reset covers every public body, including
# entities with no writable state patch and unowned world rows.
self._staged_body_wrench[rows.astype(np.intp, copy=False)] = 0.0
self._body_wrench_pending = bool(np.any(self._staged_body_wrench))
# A full default-state reset restarts the playback clock for its rows.
self._playback_time_view()[rows.astype(np.intp, copy=False)] = 0.0
def _commit_entity_reset(
self,
request: SceneResetRequest,
control_values: np.ndarray | None = None,
randomization: TensorResetRandomizationPayload | None = None,
) -> None:
super()._commit_entity_reset(request, control_values, randomization)
if self._entity_scene is None:
return
assert self._staged_body_wrench is not None
rows = np.asarray(request.env_ids, dtype=np.intp)
body_id_groups = tuple(
np.asarray(self.get_scene_layout().get_entity(patch.entity).body_ids, dtype=np.intp)
for patch in request.patches
)
body_ids = np.concatenate(body_id_groups) if body_id_groups else np.empty(0, dtype=np.intp)
self._staged_body_wrench[np.ix_(rows, body_ids)] = 0.0
self._body_wrench_pending = bool(np.any(self._staged_body_wrench))
def _worker_entrypoint(self) -> Path:
return _WORKER_PATH
def _resolve_worker_runtime(self) -> Any:
return resolve_isaacsim_runtime()
def _build_worker_environment(self, runtime: Any) -> dict[str, str]:
return build_worker_env(runtime)
def _runtime_payload(self, runtime: Any) -> dict[str, str]:
# The worker receives the resolved interpreter path as a diagnostic
# and a stable contract field. It does not import the host package.
return {
"isaacsim_python": str(runtime.python),
"isaaclab_source": (
"" if runtime.isaaclab_source is None else str(runtime.isaaclab_source)
),
}
[docs]
def get_tensor_execution(self) -> TensorExecution:
if not self._tensor_cuda_ipc_requested or self._entity_scene is None:
return TensorExecution.UNSUPPORTED
return TensorExecution.DEVICE_RESIDENT
[docs]
def get_tensor_capabilities(self) -> TensorLifecycleCapabilities:
if self.get_tensor_execution() is TensorExecution.UNSUPPORTED:
return super().get_tensor_capabilities()
return TensorLifecycleCapabilities(
execution=TensorExecution.DEVICE_RESIDENT,
state_views=True,
state_fields=frozenset({"qpos", "qvel"}),
sensor_views=True,
stepping=True,
process_topology=TensorProcessTopology.EXTERNAL_WORKER,
data_plane=TensorDataPlane.CUDA_IPC,
stream_event_ownership=(
"host records control/reset; worker records state; host consumer stream waits"
),
selected_reset=True,
selected_reset_publication=SelectedResetPublication.AUTHORITATIVE_VIEWS,
reset_randomization=True,
device_reset_randomization=True,
packed_reset=True,
packed_reset_fields=frozenset({RESET_TERM_KP, RESET_TERM_KD}),
mask_selected_reset=True,
torch_devices=("cuda",),
tracked_body_views=True,
)
[docs]
def get_public_state_widths(self) -> PublicStateWidths:
"""Return the mapped public qpos/qvel arena widths."""
if self.get_tensor_execution() is TensorExecution.UNSUPPORTED:
raise NotImplementedError(
"isaacsim public tensor state widths require tensor_cuda_ipc=true and a "
"mapped Manager-Based scene"
)
layout = self._require_mapped_entity_scene().layout
return PublicStateWidths(nq=layout.nq, nv=layout.nv)
def _cuda_tensor_reset_wire_payload(
self,
randomization: TensorResetRandomizationPayload | None,
*,
row_count: int,
device: Any,
) -> dict[str, Any] | None:
"""Validate PD gains and cross one explicit reset property boundary."""
if randomization is None or randomization.is_empty():
return None
if not isinstance(randomization, TensorResetRandomizationPayload):
raise TypeError(
"isaacsim CUDA IPC reset randomization requires a TensorResetRandomizationPayload"
)
supported = {RESET_TERM_KP, RESET_TERM_KD}
unsupported = set(randomization.requested_terms()) - supported
if unsupported:
raise NotImplementedError(
"isaacsim CUDA IPC reset randomization does not support terms: "
+ ", ".join(sorted(unsupported))
)
torch = import_torch()
wire: dict[str, Any] = {}
for term in supported:
value = getattr(randomization, term)
if value is None:
continue
require_cuda_tensor_of_dtype(value, rank=2, name=f"reset {term}", dtype="float32")
expected = (row_count, self.num_actuators)
if tuple(value.shape) != expected:
raise ValueError(
f"isaacsim CUDA IPC reset {term} must have shape {expected}, "
f"got {tuple(value.shape)}"
)
if value.device != torch.device(device):
raise ValueError(f"isaacsim CUDA IPC reset {term} device does not match arena")
if not bool(torch.isfinite(value).all().item()):
raise ValueError(f"isaacsim CUDA IPC reset {term} values must be finite")
if bool((value < 0.0).any().item()):
raise ValueError(f"isaacsim CUDA IPC reset {term} values must be nonnegative")
wire[term] = value
return wire or None
def _ensure_cuda_ipc_arena(self) -> HostCudaIpcArena:
if not self._tensor_cuda_ipc_requested or self._entity_scene is None:
raise NotImplementedError(
"isaacsim CUDA IPC tensor lifecycle requires tensor_cuda_ipc=true and a "
"mapped Manager-Based scene"
)
if self._cuda_ipc_arena is not None:
return self._cuda_ipc_arena
if self._entity_scene.owner.variant_plan is not None:
raise NotImplementedError(
"IsaacSim CUDA IPC tensor lifecycle does not support fixed variants"
)
self._require_state("CUDA IPC tensor lifecycle")
assert self._entity_scene is not None
layout = self._entity_scene.layout
sensor_descriptors = self._cuda_sensor_descriptors()
arena = HostCudaIpcArena(
num_envs=self._num_envs,
nq=layout.nq,
nv=layout.nv,
nu=layout.nu,
nbody=layout.nbody,
device="cuda",
)
try:
reply = self._request(
ISAACSIM_CUDA_ATTACH,
arena.to_payload(sensor_descriptors),
expect=ISAACSIM_CUDA_READY,
)
if (
not isinstance(reply, dict)
or reply.get("device_uuid") != arena.device_uuid
or reply.get("layout") != arena.layout.to_wire()
):
raise IsaacSimWorkerError("IsaacSim worker returned an invalid CUDA IPC handshake")
except Exception:
arena.close()
raise
# Clamp bounds are uploaded once at the cold CUDA IPC attach boundary;
# step_tensor reuses these device tensors and performs no per-step H2D.
torch = import_torch()
self._cuda_control_bounds = (
torch.as_tensor(
np.ascontiguousarray(self._entity_scene.control_lower, dtype=np.float32),
dtype=torch.float32,
device=arena.qpos.device,
),
torch.as_tensor(
np.ascontiguousarray(self._entity_scene.control_upper, dtype=np.float32),
dtype=torch.float32,
device=arena.qpos.device,
),
)
# Keep reset validation scalars resident on the worker GPU. Python
# integers passed to ``clamp`` would otherwise upload on every reset.
self._cuda_reset_row_bounds = torch.tensor(
(0, self._num_envs - 1), dtype=torch.int64, device=arena.qpos.device
)
self._cuda_reset_selected = torch.zeros(
(self._num_envs,), dtype=torch.bool, device=arena.qpos.device
)
self._cuda_reset_true = torch.ones((), dtype=torch.bool, device=arena.qpos.device)
self._cuda_ipc_arena = arena
self._cuda_tracked_body_inventory = self._cuda_body_inventory()
self._cuda_body_ids_by_name = {
f"{entity.name}/{body_name}": int(body_id)
for entity in layout.entities
for body_name, body_id in zip(entity.body_names, entity.body_ids)
}
self._cuda_sensor_spec_names = frozenset(
descriptor["name"] for descriptor in sensor_descriptors
)
self._cuda_sensor_aliases = self._cuda_sensor_aliases_from_descriptors(sensor_descriptors)
return arena
def _cuda_sensor_aliases_from_descriptors(
self, descriptors: Sequence[dict[str, Any]]
) -> dict[str, dict[str, Any]]:
aliases: dict[str, dict[str, Any]] = {}
owners: dict[str, list[str]] = {}
for slot, source in enumerate(descriptors):
descriptor = dict(source)
descriptor["slot"] = slot
names = {descriptor["name"]}
for qualified in getattr(self, "_sensor_map", {}):
if qualified == descriptor["name"] or qualified.endswith(f"/{descriptor['name']}"):
names.add(qualified)
owner = qualified.rsplit("/", 1)[0]
owners.setdefault(descriptor["name"], []).append(owner)
for name in names:
if name in aliases:
aliases[name]["ambiguous"] = True
continue
aliases[name] = {
"slot": slot,
"descriptor": descriptor,
"ambiguous": False,
}
if len(owners.get(descriptor["name"], ())) > 1:
aliases[descriptor["name"]]["ambiguous"] = True
return aliases
def _cuda_sensor_descriptors(self) -> list[dict[str, Any]]:
descriptors: list[dict[str, Any]] = []
for name in ("pelvis_local_linvel", "torso_gyro"):
candidates = [
mapped
for qualified, mapped in getattr(self, "_sensor_map", {}).items()
if qualified == name or qualified.endswith(f"/{name}")
]
if len(candidates) != 1:
continue
mapped = candidates[0]
spec, body_id = mapped
if spec.kind not in ("local_linvel", "gyro"):
continue
descriptors.append(
{
"name": name,
"kind": str(spec.kind),
"body_id": int(body_id),
"local_pos": tuple(float(value) for value in spec.local_pos),
"local_quat": tuple(float(value) for value in spec.local_quat),
}
)
return descriptors
def _cuda_body_inventory(self) -> tuple[str, ...]:
"""Return canonical tracked-body names in frozen public entity order."""
if self.get_tensor_execution() is TensorExecution.UNSUPPORTED:
raise NotImplementedError(
"IsaacSim tracked-body views require tensor_cuda_ipc=true and a mapped "
"Manager-Based scene"
)
layout = self._require_mapped_entity_scene().layout
return tuple(
f"{entity.name}/{name}" for entity in layout.entities for name in entity.body_names
)
def _require_cuda_body_inventory(self) -> tuple[str, ...]:
inventory = getattr(self, "_cuda_tracked_body_inventory", None)
if inventory is None:
return self._cuda_body_inventory()
return tuple(inventory)
def _require_cuda_sensor_aliases(self) -> dict[str, dict[str, Any]]:
aliases = getattr(self, "_cuda_sensor_aliases", None)
if aliases is None:
return {}
return dict(aliases)
def _cuda_derived_sensor_specs(self) -> dict[str, tuple[Any, int]]:
"""Return frame sensors projectable from public CUDA body state."""
return {
name: mapped
for name, mapped in getattr(self, "_sensor_map", {}).items()
if mapped[0].kind in ("framepos", "framequat", "framezaxis")
}
def _cuda_derived_sensor_view(self, name: str, arena: Any) -> Any:
candidates = [
(qualified, mapped)
for qualified, mapped in self._cuda_derived_sensor_specs().items()
if qualified == name or qualified.endswith(f"/{name}")
]
if len(candidates) != 1:
return None
_, (spec, body_id) = candidates[0]
torch = import_torch()
body = arena.body_state[:, int(body_id)]
if spec.kind == "framepos":
local_pos = torch.as_tensor(
spec.local_pos, dtype=torch.float32, device=body.device
).expand(body.shape[0], 3)
return body[:, 0:3] + torch_quat_rotate(torch, body[:, 3:7], local_pos)
local_quat = torch.as_tensor(
spec.local_quat, dtype=torch.float32, device=body.device
).expand(body.shape[0], 4)
world_quat = torch_quat_mul(torch, body[:, 3:7], local_quat)
if spec.kind == "framequat":
return world_quat
world_z = torch.zeros_like(body[:, 0:3])
world_z[:, 2] = 1.0
return torch_quat_rotate(torch, world_quat, world_z)
def _cuda_sensor_inventory(self) -> tuple[SensorDescriptor, ...]:
scalar_aliases = sorted(
alias
for alias, record in self._require_cuda_sensor_aliases().items()
if not record["ambiguous"]
)
descriptors = tuple(SensorDescriptor(name=name, width=3) for name in scalar_aliases)
descriptors += tuple(
SensorDescriptor(name=f"{prefix}{name}", width=width)
for prefix, width in (
("track_pos_w_", 3),
("track_quat_w_", 4),
("track_linvel_w_", 3),
("track_angvel_w_", 3),
)
for name in self._require_cuda_body_inventory()
)
descriptors += tuple(
SensorDescriptor(name=name, width=spec.dim)
for name, (spec, _body_id) in sorted(self._cuda_derived_sensor_specs().items())
)
return descriptors
[docs]
def get_sensor_inventory(self) -> tuple[SensorDescriptor, ...]:
"""Return the complete CUDA IPC tensor-sensor namespace without host slots."""
arena = self._ensure_cuda_ipc_arena()
del arena
return self._cuda_sensor_inventory()
[docs]
def get_state_views(
self, fields: tuple[str, ...] | str | None = None, device: Any | None = None
) -> Mapping[str, Any]:
requested = (
("qpos", "qvel")
if fields is None
else ((fields,) if isinstance(fields, str) else tuple(fields))
)
if not set(requested) <= {"qpos", "qvel"} or len(set(requested)) != len(requested):
raise KeyError("IsaacSim CUDA IPC state views only support qpos and qvel")
arena = self._ensure_cuda_ipc_arena()
if device is not None:
validate_torch_device(
(str(arena.qpos.device),),
device,
current_device=arena.device_index,
label="IsaacSim CUDA IPC state views",
)
# READY only proves that the external worker remains healthy. The
# backend-owned state IPC event, not a CPU blocking wait, orders worker
# D2D projection before every consumer stream that reads these views.
arena.wait_state()
return {name: getattr(arena, name) for name in requested}
[docs]
def get_sensor_view(self, name: str, device: Any | None = None) -> Any:
if not isinstance(name, str) or not name:
raise KeyError("IsaacSim CUDA IPC tensor sensor name must be a non-empty string")
prefix: str | None = None
for candidate in (
"track_pos_w_",
"track_quat_w_",
"track_linvel_w_",
"track_angvel_w_",
):
if name.startswith(candidate):
prefix = candidate
break
body_id: int | None = None
arena = self._ensure_cuda_ipc_arena()
if device is not None:
validate_torch_device(
(str(arena.qpos.device),),
device,
current_device=arena.device_index,
label="IsaacSim CUDA IPC sensor views",
)
if prefix is not None:
body_name = name[len(prefix) :]
body_id = self._cuda_body_ids_by_name.get(body_name)
if body_id is None or body_id >= arena.layout.nbody:
available = ", ".join(self._require_cuda_body_inventory())
raise KeyError(
f"unknown IsaacSim CUDA IPC tracked body {body_name!r}; available "
f"bodies: {available}"
)
else:
alias = self._require_cuda_sensor_aliases().get(name)
if alias is None:
has_derived = any(
qualified == name or qualified.endswith(f"/{name}")
for qualified in self._cuda_derived_sensor_specs()
)
if not has_derived:
raise KeyError(
f"unknown IsaacSim CUDA IPC tensor sensor {name!r}; available sensors: "
+ (
", ".join(
descriptor.name for descriptor in self._cuda_sensor_inventory()
)
or "<none for this scene>"
)
)
arena.wait_state()
derived = self._cuda_derived_sensor_view(name, arena)
if derived is not None:
return derived
inventory = ", ".join(
descriptor.name for descriptor in self._cuda_sensor_inventory()
)
if not inventory:
inventory = "<none for this scene>"
raise KeyError(
f"unknown IsaacSim CUDA IPC tensor sensor {name!r}; available sensors: "
f"{inventory}"
)
if alias["ambiguous"]:
raise KeyError(
f"ambiguous IsaacSim CUDA IPC tensor sensor {name!r}; use an "
"entity-qualified name"
)
arena.wait_state()
if prefix is not None:
body_state = arena.body_state
assert body_id is not None
if prefix == "track_pos_w_":
return body_state[:, body_id, 0:3]
if prefix == "track_quat_w_":
return body_state[:, body_id, 3:7]
if prefix == "track_linvel_w_":
return body_state[:, body_id, 7:10]
return body_state[:, body_id, 10:13]
alias = self._require_cuda_sensor_aliases()[name]
return arena.sensor_state[:, alias["slot"]]
[docs]
def get_tracked_body_views(
self,
body_names: Sequence[str] | None = None,
device: Any | None = None,
) -> TrackedBodyStateViews:
"""Return four ordered tracked-body blocks from one public IPC read."""
arena = self._ensure_cuda_ipc_arena()
if device is not None:
validate_torch_device(
(str(arena.qpos.device),),
device,
current_device=arena.device_index,
label="IsaacSim CUDA IPC tracked-body views",
)
declared = self._require_cuda_body_inventory()
if body_names is None:
names = declared
else:
if isinstance(body_names, (str, bytes)):
raise TypeError("IsaacSim tracked-body view names must be a sequence of strings")
names = tuple(body_names)
if not names or any(not isinstance(name, str) or not name for name in names):
raise TypeError("IsaacSim tracked-body view names must be non-empty strings")
if len(set(names)) != len(names):
raise ValueError(f"IsaacSim tracked-body view names must be unique: {names}")
missing = [name for name in names if name not in set(declared)]
if missing:
raise ValueError(
"IsaacSim tracked-body views requested bodies missing from the mapped "
f"namespace: {missing}; available={list(declared)}"
)
arena.wait_state()
body_state = arena.body_state
ids = [self._cuda_body_ids_by_name[name] for name in names]
selected_body_state = body_state[:, ids, :]
return TrackedBodyStateViews(
body_names=names,
pos_w=selected_body_state[:, :, 0:3],
quat_w=selected_body_state[:, :, 3:7],
lin_vel_w=selected_body_state[:, :, 7:10],
ang_vel_w=selected_body_state[:, :, 10:13],
)
[docs]
def step_tensor(self, ctrl: Any, nsteps: int = 1) -> dict | None:
if isinstance(nsteps, bool) or not isinstance(nsteps, int) or nsteps <= 0:
raise ValueError(f"nsteps must be a positive integer, got {nsteps!r}")
if self._pre_step_control_fn is not None:
raise NotImplementedError(
"IsaacSim CUDA IPC tensor stepping does not support host pre-step callbacks"
)
require_cuda_tensor(ctrl, rank=2, name="control")
arena = self._ensure_cuda_ipc_arena()
expected = (self._num_envs, self.num_actuators)
if tuple(ctrl.shape) != expected:
raise ValueError(f"control must have shape {expected}, got {tuple(ctrl.shape)}")
validate_torch_device(
(str(arena.ctrl.device),),
ctrl.device,
current_device=arena.device_index,
label="IsaacSim CUDA IPC control",
)
bounds = getattr(self, "_cuda_control_bounds", None)
if bounds is not None:
torch = import_torch()
ctrl = torch.clamp(ctrl, min=bounds[0], max=bounds[1])
arena.write_control(ctrl)
arena.record_control()
payload = self._request(
ISAACSIM_CUDA_STEP,
{"nsteps": nsteps},
expect=protocol.CMD_READY,
)
arena.wait_state()
timing = dict(payload.get("timing", {})) if isinstance(payload, dict) else {}
timing["cuda_ipc"] = True
return {"timing": timing}
[docs]
def set_state_tensor(
self,
env_indices: Any,
qpos: Any,
qvel: Any,
randomization: TensorResetRandomizationPayload | None = None,
) -> dict | None:
if self.get_tensor_execution() is TensorExecution.UNSUPPORTED:
return super().set_state_tensor(env_indices, qpos, qvel, randomization)
if randomization is not None and not isinstance(
randomization, TensorResetRandomizationPayload
):
raise TypeError(
"isaacsim CUDA IPC reset randomization requires a TensorResetRandomizationPayload"
)
require_cuda_tensor_of_dtype(env_indices, rank=1, name="reset env_indices", dtype="int64")
require_cuda_tensor(qpos, rank=2, name="reset qpos")
require_cuda_tensor(qvel, rank=2, name="reset qvel")
arena = self._ensure_cuda_ipc_arena()
device = arena.reset_qpos.device
for tensor, label in (
(env_indices, "IsaacSim CUDA IPC reset rows"),
(qpos, "IsaacSim CUDA IPC reset qpos"),
(qvel, "IsaacSim CUDA IPC reset qvel"),
):
validate_torch_device(
(str(device),),
tensor.device,
current_device=arena.device_index,
label=label,
)
count = int(env_indices.shape[0])
if count > self._num_envs:
raise ValueError(f"reset row count may be at most {self._num_envs}, got {count}")
expected_qpos = (count, arena.layout.nq)
expected_qvel = (count, arena.layout.nv)
if tuple(qpos.shape) != expected_qpos:
raise ValueError(f"reset qpos must have shape {expected_qpos}, got {tuple(qpos.shape)}")
if tuple(qvel.shape) != expected_qvel:
raise ValueError(f"reset qvel must have shape {expected_qvel}, got {tuple(qvel.shape)}")
if count == 0:
return {"timing": {"cuda_ipc": True, "cuda_ipc_reset_bytes": 0.0}}
torch = import_torch()
self._cuda_reset_selected.zero_()
safe_rows = env_indices.clamp(
min=self._cuda_reset_row_bounds[0], max=self._cuda_reset_row_bounds[1]
)
self._cuda_reset_selected[safe_rows] = self._cuda_reset_true
# One bounded synchronization combines row-range and uniqueness closure.
# Finite checks remain producer/task responsibility; no bulk state is
# reduced and downloaded here.
checks = torch.stack(
(env_indices.min(), env_indices.max(), self._cuda_reset_selected.sum())
)
row_min, row_max, unique_count = checks.tolist()
if row_min < 0 or row_max >= self._num_envs:
raise IndexError("IsaacSim CUDA IPC reset rows are out of range")
if unique_count != count:
raise ValueError("IsaacSim CUDA IPC reset rows must be unique")
wire_randomization = self._cuda_tensor_reset_wire_payload(
randomization, row_count=count, device=device
)
sequence = self._cuda_reset_sequence + 1
arena.write_reset(env_indices, qpos, qvel, wire_randomization)
arena.record_reset()
payload = self._request(
ISAACSIM_CUDA_RESET,
{
"count": count,
"sequence": sequence,
"has_randomization": wire_randomization is not None,
"randomization_terms": ()
if wire_randomization is None
else tuple(wire_randomization),
},
expect=protocol.CMD_READY,
)
self._cuda_reset_sequence = sequence
arena.wait_state()
timing = dict(payload.get("timing", {})) if isinstance(payload, dict) else {}
timing.update({"cuda_ipc": True, "cuda_ipc_reset_bytes": 0.0})
return {"timing": timing}
[docs]
def set_state_tensor_mask(
self,
reset_mask: Any,
qpos: Any,
qvel: Any,
randomization: TensorResetRandomizationPayload | None = None,
*,
full_refresh: bool = False,
) -> dict | None:
"""Commit full-width state through a backend-owned mask projection."""
import_torch()
if full_refresh:
raise NotImplementedError("isaacsim CUDA IPC mask reset is a selected-row boundary")
if self.get_tensor_execution() is TensorExecution.UNSUPPORTED:
raise NotImplementedError("isaacsim CUDA IPC tensor lifecycle requires a mapped scene")
require_cuda_tensor_of_dtype(reset_mask, rank=1, name="reset mask", dtype="bool")
require_cuda_tensor(qpos, rank=2, name="reset qpos")
require_cuda_tensor(qvel, rank=2, name="reset qvel")
if tuple(reset_mask.shape) != (self._num_envs,):
raise ValueError(f"isaacsim CUDA IPC reset mask must have shape {(self._num_envs,)}")
widths = self.get_public_state_widths()
expected_qpos = (self._num_envs, int(widths.nq))
expected_qvel = (self._num_envs, int(widths.nv))
if tuple(qpos.shape) != expected_qpos:
raise ValueError(f"isaacsim CUDA IPC reset qpos must have shape {expected_qpos}")
if tuple(qvel.shape) != expected_qvel:
raise ValueError(f"isaacsim CUDA IPC reset qvel must have shape {expected_qvel}")
if qpos.device != reset_mask.device or qvel.device != reset_mask.device:
raise ValueError("isaacsim CUDA IPC mask reset tensors must share one device")
validated_randomization = self._cuda_tensor_reset_wire_payload(
randomization, row_count=self._num_envs, device=reset_mask.device
)
del validated_randomization
if not bool(reset_mask.any().item()):
return {"timing": {"cuda_ipc": True, "cuda_ipc_reset_bytes": 0.0}}
rows = reset_mask.nonzero(as_tuple=False).flatten()
selected_payload: TensorResetRandomizationPayload | None = None
if randomization is not None and not randomization.is_empty():
selected_payload = TensorResetRandomizationPayload(
**{
term: getattr(randomization, term).index_select(0, rows).contiguous()
for term in (RESET_TERM_KP, RESET_TERM_KD)
if getattr(randomization, term) is not None
}
)
return self.set_state_tensor(
rows,
qpos.index_select(0, rows).contiguous(),
qvel.index_select(0, rows).contiguous(),
selected_payload,
)
[docs]
def bind_packed_reset(self, *, device: Any | None = None) -> Any:
"""Begin a packed transaction over the external CUDA IPC reset boundary."""
from unisim.backend.packed_reset import PackedResetTransaction
return PackedResetTransaction(self, device=device)
[docs]
def close(self) -> None:
arena = getattr(self, "_cuda_ipc_arena", None)
if arena is not None:
# Keep the arena reachable when caller-held views make close fail
# closed. Once its views are released, a backend ``close()`` retry
# can release the same arena instead of leaking the CUDA mapping.
arena.close()
self._cuda_ipc_arena = None
self._cuda_reset_row_bounds = None
self._cuda_reset_selected = None
self._cuda_reset_true = None
super().close()
def _mapped_contact_force_sensor_count(self) -> int:
if self._entity_scene is None:
return 0
return len(self._contact_force_sensor_payload())
def _contact_force_sensor_payload(self) -> list[dict[str, str]]:
"""Translate canonical pair sensors into worker body declarations."""
if self._entity_scene is None:
return []
metadata = self._get_scene_metadata()
body_owners = self._materialized_body_owners()
specs = [
spec
for spec in metadata.sensors.values()
if spec.kind == KIND_CONTACT_FORCE and spec.target_body_name is not None
]
records: list[dict[str, str]] = []
body_pairs: dict[tuple[str, str, str, str], str] = {}
for spec in specs:
if spec.sensor_index is None:
raise self._worker_error(
"malformed IsaacSim contact-force sensor declaration: " + spec.name
)
source = body_owners.get(spec.body_name)
target = body_owners.get(spec.target_body_name or "")
if source is None or target is None:
raise self._worker_error(
"IsaacSim collision-pair contact sensors require both bodies to belong "
f"to materialized entities (sensor {spec.name!r})"
)
body_pair = (source[0], source[1], target[0], target[1])
previous_name = body_pairs.get(body_pair)
if previous_name is not None:
raise self._worker_error(
"IsaacSim collision-pair contact sensors collapse to the same "
f"rigid-body pair: {previous_name!r} and {spec.name!r}"
)
body_pairs[body_pair] = spec.name
records.append(
{
"name": spec.name,
"source_entity": body_pair[0],
"source_body": body_pair[1],
"target_entity": body_pair[2],
"target_body": body_pair[3],
}
)
if [spec.sensor_index for spec in specs] != list(range(len(specs))):
raise self._worker_error("IsaacSim contact-force sensor row indexes are malformed")
return records
def _materialized_body_owners(self) -> dict[str, tuple[str, str]]:
assert self._entity_scene is not None
return {
entity.name + "/" + body_name: (entity.name, body_name)
for entity in self._entity_scene.layout.entities
for body_name in entity.body_names
}
def _body_net_contact_entity_payload(self) -> list[str]:
"""List entities whose bodies need the worker's per-body net force view.
Body-net ``found`` and wildcard body-net ``force`` declarations are
served from the shared per-body ``contact_force`` slot, so the worker
must report net contact forces for the entities owning the referenced
bodies. Declarations referencing unowned bodies are skipped here and
fail closed later in ``_resolve_sensor_map``.
"""
if self._entity_scene is None:
return []
metadata = self._get_scene_metadata()
body_owners = self._materialized_body_owners()
entities = set()
for spec in metadata.sensors.values():
body_net = spec.kind == KIND_CONTACT_FOUND or (
spec.kind == KIND_CONTACT_FORCE and spec.target_body_name is None
)
if not body_net:
continue
owner = body_owners.get(spec.body_name)
if owner is not None:
entities.add(owner[0])
return sorted(entities)
def _resolve_sensor_map(self) -> dict[str, tuple[Any, int]]:
"""Resolve only sensors backed by a real IsaacSim state quantity.
Mapped scenes serve collision-pair ``force`` sensors through the
dedicated PhysX reporter slot and body-net ``found``/wildcard ``force``
sensors through the per-body net-force view. The model-file worker has no
PhysX contact reporter at all, so every contact declaration fails
closed there.
"""
resolved = super()._resolve_sensor_map()
if self._entity_scene is not None:
return resolved
metadata = self._get_scene_metadata()
for name, (spec, _body_id) in tuple(resolved.items()):
if spec.kind not in (KIND_CONTACT_FOUND, KIND_CONTACT_FORCE):
continue
metadata.unsupported_sensors[name] = UnsupportedSensorSpec(
name=name,
reason=(
"IsaacSim serves contact declarations only on mapped scenes; "
"the model-file worker has no PhysX per-body or pair contact reporter"
),
)
del resolved[name]
return resolved
def _host_read_sensor(self, name: str) -> np.ndarray:
mapped = self._sensor_map.get(name)
if (
mapped is not None
and mapped[0].kind == KIND_CONTACT_FORCE
and mapped[0].target_body_name is not None
):
self._require_state("get_sensor_data")
spec: SceneSensorSpec = mapped[0]
if spec.sensor_index is None:
raise self._worker_error("IsaacSim contact-force sensor row is unresolved")
return self._slots["contact_sensor_force"][:, spec.sensor_index, :].copy()
return super()._host_read_sensor(name)
def _bind_sensor_data_reader(self, names: tuple[str, ...]) -> Callable[[], np.ndarray]:
"""Bind the cold-path NumPy shape contract over CUDA sensor views."""
def read() -> np.ndarray:
return np.concatenate(
[self.get_sensor_view(name).detach().cpu().numpy() for name in names],
axis=1,
)
return read
def _bind_model_metadata(self, meta: dict[str, Any]) -> None:
"""Validate the worker's private clone, collision, and render contract."""
expected_render_mode = self._resolved_render_mode
if expected_render_mode is None:
raise self._worker_error(
"isaacsim worker metadata arrived before the host resolved its render mode"
)
required_render_fields = {
"graphics_enabled",
"render_mode",
"render_width",
"render_height",
}
missing_render_fields = sorted(required_render_fields.difference(meta))
if missing_render_fields:
raise self._worker_error(
"isaacsim worker metadata is missing render startup fields: "
+ ", ".join(missing_render_fields)
)
reported_render_mode = meta["render_mode"]
if reported_render_mode != expected_render_mode:
raise self._worker_error(
"isaacsim worker render_mode does not match the host INIT request: "
f"worker={reported_render_mode!r}, host={expected_render_mode!r}"
)
raw_width = meta["render_width"]
raw_height = meta["render_height"]
if (
isinstance(raw_width, bool)
or not isinstance(raw_width, (int, np.integer))
or isinstance(raw_height, bool)
or not isinstance(raw_height, (int, np.integer))
or (int(raw_width), int(raw_height)) != (self._render_width, self._render_height)
):
raise self._worker_error(
"isaacsim worker render dimensions do not match the host INIT request: "
f"worker={raw_width!r}x{raw_height!r}, "
f"host={self._render_width}x{self._render_height}"
)
reported_graphics = meta["graphics_enabled"]
expected_graphics = expected_render_mode != "none"
if (
not isinstance(reported_graphics, (bool, np.bool_))
or bool(reported_graphics) != expected_graphics
):
raise self._worker_error(
"isaacsim worker graphics_enabled does not match its startup render mode: "
f"render_mode={expected_render_mode!r}, "
f"graphics_enabled={reported_graphics!r}, expected={expected_graphics!r}"
)
raw_origins = meta.get("env_origins")
if raw_origins is None:
raise self._worker_error(
"isaacsim worker did not report environment origins; refusing to expose "
"world-space clone state through the local-frame SimBackend contract"
)
origins = np.asarray(raw_origins, dtype=np.float32)
expected = (self._num_envs, 3)
if origins.shape != expected or not np.isfinite(origins).all():
raise self._worker_error(
f"isaacsim worker environment origins have invalid shape or values: "
f"got shape {origins.shape}, expected {expected}"
)
if self._num_envs > 1:
unique = np.unique(origins, axis=0)
if unique.shape[0] != self._num_envs:
raise self._worker_error(
"isaacsim worker returned duplicate environment origins; cloned actors "
"would overlap in world space"
)
if not bool(meta.get("collision_filtering_applied", False)):
raise self._worker_error(
"isaacsim worker did not apply PhysX collision filtering between environments"
)
self._validate_solver_readback(meta)
self._validate_friction_material_sharing_report(meta)
if self._entity_scene is not None:
self._validate_environment_spacing_report(meta)
self._worker_materialization_report = {
key: meta[key]
for key in ("raw_usd_cache", "role_usd_cache", "init_telemetry")
if key in meta
}
if self._entity_scene is not None:
self._validate_worker_materialization_provenance()
super()._bind_model_metadata(meta)
if self._entity_scene is not None:
self._native_entity_table("body_mass")
self._native_entity_table("body_com", width=3)
self._validated_native_geometry_records()
self._validate_reported_entity_self_collision(meta)
self._validate_reported_entity_gravity_disabled(meta)
self._validate_reported_entity_joint_velocity_limits(meta)
def _validate_reported_entity_joint_velocity_limits(self, meta: dict[str, Any]) -> None:
"""Strictly compare reported per-entity velocity limits with the INIT request."""
scene = self._require_mapped_entity_scene()
envelope = meta.get("configuration_report")
effective = envelope.get("effective") if isinstance(envelope, dict) else None
reported = (
effective.get("entity_joint_velocity_limits") if isinstance(effective, dict) else None
)
expected = {
entry["name"]: entry["joint_velocity_limits"]
for entry in scene.payload["scene_entities"]
if entry.get("joint_velocity_limits") is not None
}
if not isinstance(reported, dict) or set(reported) != set(expected):
raise self._worker_error(
"isaacsim worker did not report per-entity joint velocity limits for the "
f"mapped scene: worker={reported!r}, host={expected!r}"
)
mismatched = sorted(
name
for name, wanted in expected.items()
if not isinstance(reported[name], dict)
or set(reported[name]) != set(wanted)
or any(
not isinstance(reported[name][joint], (int, float))
or isinstance(reported[name][joint], bool)
or not np.isfinite(float(reported[name][joint]))
or float(reported[name][joint]) != float(limit)
for joint, limit in wanted.items()
)
)
if mismatched:
raise self._worker_error(
"isaacsim worker per-entity joint velocity limits do not match the host "
f"INIT request for entities: {', '.join(mismatched)}"
)
def _validate_worker_materialization_provenance(self) -> None:
"""Require importer/cache provenance for every mapped USD scene."""
report = self._worker_materialization_report
if not isinstance(report, dict):
raise self._worker_error("isaacsim worker omitted USD materialization provenance")
required_versions = {
"isaacsim",
"isaaclab",
"isaacsim.asset.importer.mjcf",
"python",
}
required_entry = {
"identity",
"source_digest",
"entity",
"variant",
"hit",
"materialize_ms",
"artifact_files",
"artifact_bytes",
"runtime_versions",
}
for kind in ("raw_usd_cache", "role_usd_cache"):
cache = report.get(kind)
entries = cache.get("entries") if isinstance(cache, dict) else None
if not isinstance(entries, (list, tuple)) or not entries:
raise self._worker_error(
"isaacsim worker omitted " + kind + " provenance for the mapped scene"
)
for entry in entries:
versions = entry.get("runtime_versions") if isinstance(entry, dict) else None
if (
not isinstance(entry, dict)
or not required_entry.issubset(entry)
or not isinstance(versions, dict)
or not required_versions.issubset(versions)
):
raise self._worker_error(
"isaacsim worker returned malformed " + kind + " provenance"
)
def _validate_reported_entity_gravity_disabled(self, meta: dict[str, Any]) -> None:
"""Strictly compare reported per-entity gravity state with the INIT request."""
scene = self._require_mapped_entity_scene()
envelope = meta.get("configuration_report")
effective = envelope.get("effective") if isinstance(envelope, dict) else None
reported = effective.get("entity_gravity_disabled") if isinstance(effective, dict) else None
expected = {
entry["name"]: bool(entry["gravity_disabled"])
for entry in scene.payload["scene_entities"]
}
if not isinstance(reported, dict) or set(reported) != set(expected):
raise self._worker_error(
"isaacsim worker did not report per-entity gravity state for the "
f"mapped scene: worker={reported!r}, host={expected!r}"
)
mismatched = sorted(
name
for name, wanted in expected.items()
if not isinstance(reported[name], bool) or reported[name] != wanted
)
if mismatched:
raise self._worker_error(
"isaacsim worker per-entity gravity state does not match the host INIT "
f"request for entities: {', '.join(mismatched)}"
)
def _validate_reported_entity_self_collision(self, meta: dict[str, Any]) -> None:
"""Strictly compare reported per-entity self-collision with the INIT request."""
scene = self._require_mapped_entity_scene()
envelope = meta.get("configuration_report")
effective = envelope.get("effective") if isinstance(envelope, dict) else None
collision_filter = (
effective.get("collision_filter") if isinstance(effective, dict) else None
)
reported = (
collision_filter.get("self_collision") if isinstance(collision_filter, dict) else None
)
expected = {
entry["name"]: bool(entry["self_collision"])
for entry in scene.payload["scene_entities"]
}
if not isinstance(reported, dict) or set(reported) != set(expected):
raise self._worker_error(
"isaacsim worker did not report per-entity self-collision for the "
f"mapped scene: worker={reported!r}, host={expected!r}"
)
mismatched = sorted(
name
for name, wanted in expected.items()
if not isinstance(reported[name], bool) or reported[name] != wanted
)
if mismatched:
raise self._worker_error(
"isaacsim worker self-collision does not match the host INIT request "
f"for entities: {', '.join(mismatched)}"
)
def _validate_solver_readback(self, meta: dict[str, Any]) -> None:
"""Fail closed when the worker's engine readback misses the INIT request."""
requested = self._physx_solver.to_payload()
if not requested:
return
envelope = meta.get("configuration_report")
if not isinstance(envelope, dict) or not isinstance(envelope.get("effective"), dict):
raise self._worker_error(
"isaacsim worker omitted its configuration report although PhysX solver "
"overrides were requested"
)
effective = envelope["effective"]
raw_readback = envelope.get("engine_readback")
readback_fields = set(raw_readback) if isinstance(raw_readback, (list, tuple)) else set()
for field, value in requested.items():
# The GPU buffer capacities are PhysX carb settings with no USD
# attribute; the worker reports the authored value instead of an
# engine readback, and the strict comparison below still applies.
if field not in PHYSX_SOLVER_AUTHORED_FIELDS and field not in readback_fields:
raise self._worker_error(
f"isaacsim worker did not read back the requested PhysX setting "
f"{field!r} from the engine"
)
if field not in effective:
raise self._worker_error(
"isaacsim worker configuration report is missing the requested "
f"PhysX setting {field!r}"
)
if not solver_value_matches(field, value, effective[field]):
raise self._worker_error(
f"isaacsim worker PhysX {field} does not match the host INIT request: "
f"worker={effective[field]!r}, host={value!r}"
)
def _validate_friction_material_sharing_report(self, meta: dict[str, Any]) -> None:
"""Fail closed when the worker did not honor the sharing opt-in."""
expected = getattr(self, "_share_friction_materials", False)
if not expected and self._entity_scene is None:
return
envelope = meta.get("configuration_report")
effective = envelope.get("effective") if isinstance(envelope, dict) else None
reported = (
effective.get("share_friction_materials") if isinstance(effective, dict) else None
)
if not isinstance(reported, bool) or reported != expected:
raise self._worker_error(
"isaacsim friction material sharing does not match the host INIT request: "
f"worker={reported!r}, host={expected!r}"
)
def _validate_environment_spacing_report(self, meta: dict[str, Any]) -> None:
"""Fail closed when GridCloner did not honor mapped-scene spacing."""
if getattr(self, "_entity_scene", None) is None:
return
envelope = meta.get("configuration_report")
effective = envelope.get("effective") if isinstance(envelope, dict) else None
reported = effective.get("env_spacing") if isinstance(effective, dict) else None
expected = getattr(self, "_env_spacing", 2.0)
if isinstance(reported, bool) or not isinstance(reported, (int, float)):
raise self._worker_error(
f"isaacsim worker reported invalid environment spacing: {reported!r}"
)
if not bool(np.isclose(reported, expected, rtol=0.0, atol=1e-6)):
raise self._worker_error(
"isaacsim environment spacing does not match the host INIT request: "
f"worker={reported!r}, host={expected!r}"
)
raw_origins = meta.get("env_origins")
if self._num_envs > 1 and raw_origins is not None:
origins = np.asarray(raw_origins, dtype=np.float32)
tolerance = max(1e-6, expected * 1e-6)
axis_steps: list[np.ndarray] = []
for axis in range(2):
coordinates = np.unique(origins[:, axis])
steps = np.diff(coordinates)
axis_steps.append(steps[steps > tolerance])
observed = (
np.concatenate(axis_steps)
if any(value.size for value in axis_steps)
else np.array([])
)
if observed.size == 0 or not np.allclose(observed, expected, rtol=0.0, atol=tolerance):
raise self._worker_error(
"isaacsim environment origins do not implement the requested "
f"GridCloner spacing: expected={expected!r}, "
f"minimum_origin_step={float(np.min(observed)) if observed.size else None!r}"
)
def _require_mapped_entity_scene(self) -> PreparedWorkerScene:
if self._entity_scene is None:
raise NotImplementedError(
"IsaacSim body-property readback requires an explicit entity scene"
)
return self._entity_scene
def _consume_entity_reset_randomization(self, response: Any) -> None:
"""Replace current native property records from the reset barrier."""
scene = self._require_mapped_entity_scene()
raw_records = response.get("native_entity_records") if isinstance(response, dict) else None
if not isinstance(raw_records, list):
raise self._worker_error(
"isaacsim worker omitted native records after property mutation"
)
records = {record.get("name"): record for record in raw_records if isinstance(record, dict)}
if len(records) != len(raw_records) or set(records) != set(self.get_entity_names()):
raise self._worker_error(
"isaacsim worker property-mutation records do not match the frozen entity layout"
)
for entity in scene.layout.entities:
current = self._native_entity_records.get(entity.name)
if current is None:
raise self._worker_error("worker omitted native entity properties: " + entity.name)
reported = records[entity.name]
for field in ("body_mass", "body_com", "body_inertia", "geom_friction"):
if field not in reported:
raise self._worker_error(
f"worker omitted native {field} after property mutation: {entity.name}"
)
if field in current:
current[field] = reported[field]
# Validate the accepted records before returning from the reset barrier.
self._native_entity_table("body_mass")
self._native_entity_table("body_com", width=3)
self._validate_native_body_inertia_records()
self._validated_native_geometry_records()
def _validate_native_body_inertia_records(self) -> None:
"""Validate full native inertia matrices before a reset barrier returns."""
scene = self._require_mapped_entity_scene()
self._require_materialized()
for entity in scene.layout.entities:
record = self._native_entity_records.get(entity.name)
if record is None:
raise self._worker_error("worker omitted native entity properties: " + entity.name)
expected = (self._num_envs, len(entity.body_ids), 3, 3)
try:
values = np.asarray(record.get("body_inertia"), dtype=np.float32)
except (TypeError, ValueError) as exc:
raise self._worker_error(
f"worker native body_inertia is malformed for entity {entity.name}: "
f"expected shape {expected}"
) from exc
if values.shape != expected or not np.isfinite(values).all():
raise self._worker_error(
f"worker native body_inertia is malformed for entity {entity.name}: "
f"got shape {values.shape}, expected {expected}"
)
def _canonical_body_table(self, field: str, width: int | None = None) -> np.ndarray:
scene = self._require_mapped_entity_scene()
vector = width == 3 or field in ("body_com", "body_ipos")
model_field = "body_ipos" if field == "body_com" else field
expected = (scene.layout.nbody, 3) if vector else (scene.layout.nbody,)
try:
canonical = np.asarray(getattr(scene.owner.model, model_field), dtype=np.float32)
except (TypeError, ValueError) as exc:
raise self._worker_error(
f"compiled canonical {field} is malformed: expected shape {expected}"
) from exc
if canonical.shape != expected or not np.isfinite(canonical).all():
raise self._worker_error(
f"compiled canonical {field} is malformed: got shape {canonical.shape}, "
f"expected {expected}"
)
return canonical.copy()
def _native_entity_table(self, field: str, width: int | None = None) -> np.ndarray:
scene = self._require_mapped_entity_scene()
self._require_materialized()
layout = scene.layout
shape = (
(self._num_envs, layout.nbody)
if width is None
else (self._num_envs, layout.nbody, width)
)
canonical = self._canonical_body_table(field)
values = np.broadcast_to(canonical, shape).copy()
for entity in layout.entities:
record = self._native_entity_records.get(entity.name)
if record is None:
raise self._worker_error("worker omitted native entity properties: " + entity.name)
entity_shape = (
(self._num_envs, len(entity.body_ids))
if width is None
else (self._num_envs, len(entity.body_ids), width)
)
try:
raw = np.asarray(record.get(field), dtype=np.float32)
except (TypeError, ValueError) as exc:
raise self._worker_error(
f"worker native {field} is malformed for entity {entity.name}: "
f"expected shape {entity_shape}"
) from exc
if raw.shape != entity_shape or not np.isfinite(raw).all():
raise self._worker_error(
f"worker native {field} is malformed for entity {entity.name}: "
f"got shape {raw.shape}, expected {entity_shape}"
)
if width is None:
values[:, entity.body_ids] = raw
else:
values[:, entity.body_ids, :] = raw
return values
[docs]
def get_body_mass(self) -> np.ndarray:
"""Return native per-environment body masses in public body-id order."""
return self._native_entity_table("body_mass")
[docs]
def get_body_ipos(self, env_ids: Sequence[int] | np.ndarray | None = None) -> np.ndarray:
"""Return canonical defaults or native selected materialization COM offsets."""
self._require_mapped_entity_scene()
if env_ids is None:
return self._canonical_body_table("body_ipos")
ids = self._validate_env_ids(env_ids)
return self._native_entity_table("body_com", width=3)[ids]
@staticmethod
def _validate_native_geometry_audit_row(row: Any, entity_name: str, env_index: int) -> None:
"""Validate one serializable cold-path USD collision audit row."""
label = f"geometry audit for entity {entity_name!r} env {env_index}"
if not isinstance(row, dict) or set(row) != _NATIVE_GEOMETRY_AUDIT_FIELDS:
raise IsaacSimWorkerError("worker malformed " + label)
if (
row["schema_version"] != 1
or not isinstance(row["path"], str)
or not row["path"]
or not isinstance(row["usd_authored_type_name"], str)
or not isinstance(row["usd_schema_type_name"], str)
or not (row["usd_authored_type_name"] or row["usd_schema_type_name"])
or not isinstance(row["shape_kind"], str)
or not row["shape_kind"]
or not isinstance(row["collision_enabled"], bool)
or not isinstance(row["physx_collision_api_present"], bool)
or not (
row["mesh_approximation"] is None
or (isinstance(row["mesh_approximation"], str) and row["mesh_approximation"])
)
):
raise IsaacSimWorkerError("worker malformed " + label)
for prefix in ("contact_offset", "rest_offset"):
authored = row[prefix + "_authored"]
schema_value = row[prefix + "_schema_value"]
simulation_determined = row[prefix + "_simulation_determined"]
api_present = row["physx_collision_api_present"]
if (
(api_present and not isinstance(simulation_determined, bool))
or (
not api_present and not (simulation_determined is None and schema_value is None)
)
or (
authored is not None
and (
isinstance(authored, bool)
or not isinstance(authored, (int, float))
or not np.isfinite(authored)
)
)
or (
api_present
and simulation_determined
and schema_value != "simulation_determined"
)
or (
api_present
and not simulation_determined
and (
isinstance(schema_value, bool)
or not isinstance(schema_value, (int, float))
or not np.isfinite(schema_value)
)
)
):
raise IsaacSimWorkerError("worker malformed " + label + " " + prefix)
for field in ("local_transform", "world_transform"):
try:
matrix = np.asarray(row[field], dtype=np.float64)
except (TypeError, ValueError) as exc:
raise IsaacSimWorkerError("worker malformed " + label + " " + field) from exc
if (
matrix.shape != (4, 4)
or not np.isfinite(matrix).all()
or not np.allclose(matrix[:, 3], (0, 0, 0, 1), rtol=0.0, atol=1e-6)
):
raise IsaacSimWorkerError("worker malformed " + label + " " + field)
for scope in ("local", "world"):
empty = row[scope + "_bounds_empty"]
minimum = row[scope + "_bounds_min"]
maximum = row[scope + "_bounds_max"]
extent = row[scope + "_extent"]
if not isinstance(empty, bool) or (
(minimum is None) != empty
or (maximum is None) != empty
or (extent is None) != empty
):
raise IsaacSimWorkerError("worker malformed " + label + " bounds")
if empty:
continue
try:
values_min = np.asarray(minimum, dtype=np.float64)
values_max = np.asarray(maximum, dtype=np.float64)
values_extent = np.asarray(extent, dtype=np.float64)
except (TypeError, ValueError) as exc:
raise IsaacSimWorkerError("worker malformed " + label + " bounds") from exc
if (
values_min.shape != (3,)
or values_max.shape != (3,)
or values_extent.shape != (3,)
or not np.isfinite(values_min).all()
or not np.isfinite(values_max).all()
or not np.isfinite(values_extent).all()
or np.any(values_max < values_min)
or not np.allclose(values_extent, values_max - values_min, rtol=0.0, atol=1e-6)
):
raise IsaacSimWorkerError("worker malformed " + label + " bounds")
tolerance = row["extent_tolerance"]
if (
not isinstance(row["zero_extent"], bool)
or not isinstance(row["zero_area"], bool)
or not isinstance(row["zero_volume_extent"], bool)
or not isinstance(row["planar_extent"], bool)
or isinstance(tolerance, bool)
or not isinstance(tolerance, (int, float))
or not np.isfinite(tolerance)
or tolerance < 0.0
):
raise IsaacSimWorkerError("worker malformed " + label + " extent conclusion")
mesh_fields = (
"mesh_vertex_count",
"mesh_unique_vertex_count",
"mesh_face_count",
"mesh_face_vertex_index_count",
"mesh_tolerance",
"mesh_all_points_coincident",
"mesh_points_min",
"mesh_points_max",
"mesh_points_extent",
"mesh_points_zero_extent",
"mesh_surface_area",
"mesh_zero_area",
)
if row["shape_kind"] != "mesh":
if any(row[field] is not None for field in mesh_fields):
raise IsaacSimWorkerError("worker malformed " + label + " mesh summary")
return
if any(
isinstance(row[field], bool) or not isinstance(row[field], int) or row[field] < 0
for field in mesh_fields[:4]
):
raise IsaacSimWorkerError("worker malformed " + label + " mesh counts")
mesh_tolerance = row["mesh_tolerance"]
surface_area = row["mesh_surface_area"]
if (
isinstance(mesh_tolerance, bool)
or not isinstance(mesh_tolerance, (int, float))
or not np.isfinite(mesh_tolerance)
or mesh_tolerance < 0.0
or isinstance(surface_area, bool)
or not isinstance(surface_area, (int, float))
or not np.isfinite(surface_area)
or surface_area < 0.0
or not isinstance(row["mesh_all_points_coincident"], bool)
or not isinstance(row["mesh_points_zero_extent"], bool)
or not isinstance(row["mesh_zero_area"], bool)
):
raise IsaacSimWorkerError("worker malformed " + label + " mesh summary")
try:
values_min = np.asarray(row["mesh_points_min"], dtype=np.float64)
values_max = np.asarray(row["mesh_points_max"], dtype=np.float64)
values_extent = np.asarray(row["mesh_points_extent"], dtype=np.float64)
except (TypeError, ValueError) as exc:
raise IsaacSimWorkerError("worker malformed " + label + " mesh points") from exc
if (
values_min.shape != (3,)
or values_max.shape != (3,)
or values_extent.shape != (3,)
or not np.isfinite(values_min).all()
or not np.isfinite(values_max).all()
or not np.isfinite(values_extent).all()
or np.any(values_max < values_min)
or not np.allclose(values_extent, values_max - values_min, rtol=0.0, atol=1e-6)
):
raise IsaacSimWorkerError("worker malformed " + label + " mesh points")
@staticmethod
def _scattered_source_geometry_intent(entity: Any, source: dict[str, Any]) -> dict[str, Any]:
"""Scatter optional-slot source rows into frozen public geom order."""
public = [(geom.name, geom.body_name) for geom in entity.geoms]
actual = list(zip(source["geom_names"], source["geom_body_names"]))
slots: list[int] = []
cursor = 0
for pair in actual:
while cursor < len(public) and public[cursor] != pair:
cursor += 1
if cursor == len(public):
raise IsaacSimWorkerError(
"geometry audit source names are outside the public layout for entity "
+ entity.name
)
slots.append(cursor)
cursor += 1
present = set(slots)
result: dict[str, Any] = {
"geom_source_types": ["absent"] * len(public),
"geom_source_sizes": [[0.0, 0.0, 0.0] for _ in public],
"geom_source_poses": [[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0] for _ in public],
"geom_contype": [0] * len(public),
"geom_conaffinity": [0] * len(public),
"geom_placeholders": [0 if slot in present else 1 for slot in range(len(public))],
}
for source_index, slot in enumerate(slots):
for field in (
"geom_source_types",
"geom_source_sizes",
"geom_source_poses",
"geom_contype",
"geom_conaffinity",
):
result[field][slot] = source[field][source_index]
return result
@classmethod
def _validate_native_geometry_audit_record(
cls, entity: Any, entry: dict[str, Any], record: dict[str, Any], num_envs: int
) -> None:
"""Validate source/native geometry rows against immutable variant intent."""
if record.get("geometry_audit_schema_version") != 1:
raise IsaacSimWorkerError(
"worker omitted or has unsupported geometry audit schema for entity " + entity.name
)
rows: dict[str, Any] = {}
extents: dict[str, Any] = {}
audits = record.get("geom_native_audits")
for field, width in (
("geom_source_types", None),
("geom_source_sizes", 3),
("geom_source_poses", 7),
("geom_native_audits", None),
):
values = record.get(field)
if (
not isinstance(values, list)
or len(values) != num_envs
or any(not isinstance(row, list) for row in values)
or any(len(row) != len(entity.geoms) for row in values)
):
raise IsaacSimWorkerError(
f"worker geometry audit field {field} is malformed for entity {entity.name}"
)
if width is None:
rows[field] = values
else:
if not entity.geoms:
if any(row != [] for row in values):
raise IsaacSimWorkerError(
f"worker geometry audit field {field} is malformed "
f"for entity {entity.name}"
)
continue
try:
extents[field] = np.asarray(values, dtype=np.float64)
except (TypeError, ValueError) as exc:
raise IsaacSimWorkerError(
f"worker geometry audit field {field} is malformed for entity {entity.name}"
) from exc
if (
extents[field].shape != (num_envs, len(entity.geoms), width)
or not np.isfinite(extents[field]).all()
):
raise IsaacSimWorkerError(
f"worker geometry audit field {field} is malformed for entity {entity.name}"
)
audit_rows: Any = audits
for env_index in range(num_envs):
if not entity.geoms:
continue
variant = int(entry["assignment"][env_index])
source = cls._scattered_source_geometry_intent(entity, entry["variants"][variant])
if rows["geom_source_types"][env_index] != source["geom_source_types"]:
raise IsaacSimWorkerError(
f"worker geometry audit source types differ for entity {entity.name} "
f"in env {env_index}"
)
for field in ("geom_source_sizes", "geom_source_poses"):
if not np.allclose(
extents[field][env_index],
np.asarray(source[field], dtype=np.float64),
rtol=1e-6,
atol=1e-7,
):
raise IsaacSimWorkerError(
f"worker geometry audit source {field} differ for entity {entity.name} "
f"in env {env_index}"
)
placeholders = source["geom_placeholders"]
for slot, audit in enumerate(audit_rows[env_index]):
source_colliding = bool(
source["geom_contype"][slot] or source["geom_conaffinity"][slot]
)
placeholder = bool(placeholders[slot])
native_expected = bool(source_colliding or placeholder)
if native_expected != (audit is not None):
raise IsaacSimWorkerError(
f"worker geometry audit native presence differs for entity "
f"{entity.name} geom {entity.geoms[slot].name!r} in env {env_index}"
)
if audit is not None:
cls._validate_native_geometry_audit_row(audit, entity.name, env_index)
expected_enabled = bool(entry["collision_enabled"] and source_colliding)
if audit["collision_enabled"] != expected_enabled:
raise IsaacSimWorkerError(
f"worker geometry audit collision state differs for entity "
f"{entity.name} geom {entity.geoms[slot].name!r} in env {env_index}"
)
def _validated_native_geometry_records(
self,
) -> tuple[tuple[Any, int, np.ndarray, np.ndarray], ...]:
"""Validate and return native geometry records in public entity order.
The returned tuple contains ``(entity, public offset, masks, friction)``.
Identity is checked against the frozen layout while values remain native
worker readback; contact state is role-immutable and friction may vary by
immutable variant assignment.
"""
scene = self._require_mapped_entity_scene()
self._require_materialized()
result = []
public_offset = 0
for entity, entry in zip(scene.layout.entities, scene.payload["scene_entities"]):
record = self._native_entity_records.get(entity.name)
if record is None:
raise self._worker_error("worker omitted native entity geometry: " + entity.name)
self._validate_native_geometry_audit_record(entity, entry, record, self._num_envs)
names = record.get("geom_names")
body_names = record.get("geom_body_names")
expected_names = [geom.name for geom in entity.geoms]
expected_bodies = [geom.body_name for geom in entity.geoms]
# uniform_public_layout permits optional mesh-geom slots to be
# absent per environment (e.g. eraser heads). The worker normally
# pads those slots with placeholder prims, so rows arrive
# full-length; per-row order-preserving subsets are still accepted
# here and padded in the native value tables as a fail-closed
# tolerance. Geom names are unique per entity, so name matching
# fixes each row's public slots exactly.
if (
not isinstance(names, list)
or len(names) != self._num_envs
or any(not isinstance(row, list) for row in names)
):
raise self._worker_error(
f"worker native geom_names are malformed for entity {entity.name}: "
f"expected {self._num_envs} rows"
)
if (
not isinstance(body_names, list)
or len(body_names) != self._num_envs
or any(not isinstance(row, list) for row in body_names)
):
raise self._worker_error(
f"worker native geom_body_names are malformed for entity {entity.name}: "
f"expected {self._num_envs} rows"
)
slot_rows: list[list[int]] = []
for row_names, row_bodies in zip(names, body_names):
slots: list[int] = []
cursor = 0
for name in row_names:
while cursor < len(expected_names) and expected_names[cursor] != name:
cursor += 1
if cursor == len(expected_names):
raise self._worker_error(
f"worker native geom_names differ from the frozen layout for "
f"entity {entity.name}: row entry {name!r} is outside the "
f"public layout"
)
slots.append(cursor)
cursor += 1
if len(row_bodies) != len(row_names):
raise self._worker_error(
f"worker native geom_body_names are malformed for entity "
f"{entity.name}: row length differs from geom_names"
)
for slot, body in zip(slots, row_bodies):
if body != expected_bodies[slot]:
raise self._worker_error(
f"worker native geom_body_names differ from the frozen layout "
f"for entity {entity.name}: geom {expected_names[slot]!r} sits "
f"on {body!r}, expected {expected_bodies[slot]!r}"
)
slot_rows.append(slots)
if not entity.geoms:
empty_rows: list[list[int]] = [[] for _ in range(self._num_envs)]
if (
record.get("geom_contact_masks") != empty_rows
or record.get("geom_friction") != empty_rows
):
raise self._worker_error(
"worker native geometry values are malformed for entity " + entity.name
)
masks = np.empty((self._num_envs, 0, 2), dtype=np.int32)
friction = np.empty((self._num_envs, 0, 3), dtype=np.float32)
else:
raw_masks = record.get("geom_contact_masks")
raw_friction = record.get("geom_friction")
masks = np.zeros((self._num_envs, len(entity.geoms), 2), dtype=np.int32)
friction = np.zeros((self._num_envs, len(entity.geoms), 3), dtype=np.float32)
try:
for env_index, slots in enumerate(slot_rows):
row_masks = np.asarray(raw_masks[env_index], dtype=np.int32)
row_friction = np.asarray(raw_friction[env_index], dtype=np.float32)
if row_masks.shape != (len(slots), 2) or row_friction.shape != (
len(slots),
3,
):
raise ValueError("native geom row shape mismatch")
for local, slot in enumerate(slots):
masks[env_index, slot] = row_masks[local]
friction[env_index, slot] = row_friction[local]
except (TypeError, ValueError, IndexError) as exc:
raise self._worker_error(
f"worker native geom_contact_masks or geom_friction is malformed "
f"for entity {entity.name}"
) from exc
if not np.isin(masks, (0, 1)).all():
raise self._worker_error(
f"worker native geom_contact_masks are malformed for entity "
f"{entity.name}: values must be 0/1"
)
if not np.isfinite(friction).all() or np.any(friction < 0.0):
raise self._worker_error(
f"worker native geom_friction is malformed for entity {entity.name}"
)
# Contact state is role-immutable: audit per-slot equality
# across the environments whose assigned variant carries the
# slot. Variants that omit an optional uniform_public_layout
# slot surface it as a disabled placeholder row (or omit the
# row entirely); a placeholder row must report a fully
# disabled mask.
variant_geoms = [
frozenset(variant.get("geom_names") or ()) for variant in entry["variants"]
]
for slot in range(len(entity.geoms)):
carrying: list[int] = []
for env_index, slots in enumerate(slot_rows):
if slot not in slots:
continue
assigned = variant_geoms[int(entry["assignment"][env_index])]
if expected_names[slot] in assigned:
carrying.append(env_index)
elif np.any(masks[env_index, slot] != 0):
raise self._worker_error(
f"worker native geom_contact_masks report an omitted "
f"optional slot as enabled for entity {entity.name}: "
f"geom {expected_names[slot]!r} in env {env_index}"
)
if carrying and not np.all(masks[carrying, slot] == masks[carrying[0], slot]):
raise self._worker_error(
"worker native geom_contact_masks vary across environments "
"for entity " + entity.name
)
result.append((entity, public_offset, masks, friction))
public_offset += len(entity.geoms)
if public_offset != scene.layout.ngeom:
raise self._worker_error("worker native geometry does not cover the frozen layout")
return tuple(result)
[docs]
def get_geom_names(self) -> tuple[str, ...]:
"""Return qualified geometry names in frozen public geometry order."""
if self._entity_scene is None:
raise NotImplementedError(f"{self.__class__.__name__} does not expose geom names")
return tuple(
entity.name + "/" + geom.name
for entity in self._entity_scene.layout.entities
for geom in entity.geoms
)
[docs]
def get_geom_body_ids(self) -> np.ndarray:
"""Return owning public body IDs in frozen geometry order."""
if self._entity_scene is None:
raise NotImplementedError(f"{self.__class__.__name__} does not expose geom body ids")
body_ids = np.empty(self._entity_scene.layout.ngeom, dtype=np.int32)
offset = 0
for entity in self._entity_scene.layout.entities:
owners = dict(zip(entity.body_names, entity.body_ids, strict=True))
for geom in entity.geoms:
body_ids[offset] = owners[geom.body_name]
offset += 1
return body_ids
[docs]
def get_geom_friction(self) -> np.ndarray:
"""Return native effective per-environment friction in public order.
Columns are IsaacSim ``[static_friction, dynamic_friction, 0]``. PhysX
has no native equivalent for MuJoCo's torsional/rolling coefficients.
"""
if self._entity_scene is None:
raise NotImplementedError(f"{self.__class__.__name__} does not expose geom friction")
scene = self._require_mapped_entity_scene()
records = self._validated_native_geometry_records()
result = np.empty((self._num_envs, scene.layout.ngeom, 3), dtype=np.float32)
for entity, offset, _masks, friction in records:
result[:, offset : offset + len(entity.geoms), :] = friction
return result
[docs]
def get_play_capabilities(self) -> PlayCapabilities:
"""Return the Kit viewer/camera plus physics-state playback capabilities.
Both scene profiles implement the playback contract. The mapped
entity scene serves worker-maintained public ``qpos``/``qvel`` rows;
the model-file whole-model profile assembles the same snapshot from its
root/dof wire slots once the construction-time joint-order validation
passes. Model-file scenes whose MJCF topology the synthetic 7/6-root wire
cannot replay (fixed base, ball joints, reordered freejoints) keep
failing closed here.
"""
return PlayCapabilities(
supports_native_interactive_renderer=True,
supports_physics_state_playback=self._active_playback_state_layout() is not None,
supports_native_video_capture=True,
)
def _active_playback_state_layout(self) -> PhysicsStateLayout | None:
"""Return the snapshot layout of the active profile, or ``None``.
``None`` means physics-state playback is unsupported for this scene
and every contract entry point fails closed through the base class.
"""
if self._entity_scene is not None:
layout = self._entity_scene.layout
return PhysicsStateLayout(nq=int(layout.nq), nv=int(layout.nv))
return self._model_file_playback_state_layout()
def _model_file_playback_state_layout(self) -> PhysicsStateLayout | None:
"""Return the cached model-file snapshot layout (``None`` stays fail-closed)."""
if self._entity_scene is not None:
return None
if not hasattr(self, "_model_file_playback_layout"):
raise RuntimeError("IsaacSim playback layout must be initialized by the constructor")
return self._model_file_playback_layout
[docs]
def get_physics_state_layout(self) -> PhysicsStateLayout:
"""Return the ``[time, qpos, qvel]`` snapshot layout (no mocap tail yet)."""
layout = self._active_playback_state_layout()
if layout is None:
return super().get_physics_state_layout()
return layout
[docs]
def get_physics_state(self) -> np.ndarray:
"""Assemble contract ``[time, qpos, qvel]`` rows from the public slots.
The worker refreshes the public shared-memory rows after every step
and reset, so the snapshot is a host-side copy of those rows plus the
leading time column; no new IPC frame or worker-side asset access is
involved. The worker protocol carries no simulation clock, so time
comes from the per-env host accumulator advanced on every public step
path and zeroed on full state resets, mirroring the mujoco backend's
time semantics.
The mapped profile copies the MJCF-order ``qpos``/``qvel`` rows
verbatim. The model-file profile projects its wire slots: ``root_state``
carries the synthetic floating root as ``[pos, quat wxyz, world lin
vel (link origin), world ang vel]`` and ``dof_state`` the scalar
joints in MJCF document order, so the snapshot's MuJoCo-convention
qvel angular block is the world angular velocity rotated into the
body frame — exactly the convention the model-file ``set_state`` path
consumes, which keeps snapshot restore a lossless round trip.
"""
layout = self._active_playback_state_layout()
if layout is None:
return super().get_physics_state()
self._require_state("physics snapshot")
out = np.empty((self._num_envs, layout.state_width), dtype=np.float32)
out[:, 0] = self._playback_time_view()
if self._entity_scene is not None:
out[:, 1 : 1 + layout.nq] = self._slots["qpos"]
out[:, 1 + layout.nq :] = self._slots["qvel"]
return out
root_state = self._slots["root_state"]
dof_state = self._slots["dof_state"]
out[:, 1:8] = root_state[:, :7]
out[:, 8 : 1 + layout.nq] = dof_state[:, :, 0]
qvel = 1 + layout.nq
out[:, qvel : qvel + 3] = root_state[:, 7:10]
out[:, qvel + 3 : qvel + 6] = protocol.quat_rotate_inverse(
root_state[:, 3:7], root_state[:, 10:13]
)
out[:, qvel + 6 :] = dof_state[:, :, 1]
return out
[docs]
def set_physics_state(self, state: np.ndarray) -> None:
"""Restore a ``get_physics_state`` snapshot through the set_state path."""
layout = self._active_playback_state_layout()
if layout is None:
return super().set_physics_state(state)
state_array = np.asarray(state, dtype=np.float32)
expected = (self._num_envs, layout.state_width)
if state_array.shape != expected:
raise ValueError(
"isaacsim physics snapshot must use [time, qpos, qvel] layout with shape "
f"{expected}, got {state_array.shape}"
)
parts = layout.split_state(state_array)
# The split columns are strided views; the shared-memory reset upload
# expects contiguous rows.
self._host_set_state(
np.arange(self._num_envs, dtype=np.intp),
np.ascontiguousarray(parts.qpos),
np.ascontiguousarray(parts.qvel),
)
# set_state restarts the playback clock for its rows; the snapshot's
# time column is authoritative for the restore.
self._playback_time_view()[...] = parts.time
[docs]
def get_playback_model(self, env_index: int | None = None) -> Any:
"""Return the MJCF playback source for detached MuJoCo replay.
The mapped profile keeps the shared implementation (fixed-variant
scenes require an explicit ``env_index``). The model-file profile
replays through its whole-MJCF construction source; the adapter
rejects model-file fixed-variant plans at construction, so one source
serves every row.
"""
if self._entity_scene is None and self._active_playback_state_layout() is not None:
if env_index is not None:
if isinstance(env_index, bool) or not isinstance(env_index, int):
raise TypeError("env_index must be an integer")
if not 0 <= env_index < self._num_envs:
raise IndexError("playback environment is out of range")
return str(Path(self._scene.model_file).expanduser())
return super().get_playback_model(env_index)
def _host_set_state(
self,
env_indices: np.ndarray,
qpos: np.ndarray,
qvel: np.ndarray,
randomization: TensorResetRandomizationPayload | None = None,
) -> dict[str, dict[str, float]]:
result = super()._host_set_state(env_indices, qpos, qvel, randomization)
# A full state write restarts the playback clock for its rows,
# matching the mujoco backend's reset-world time clear. This also
# covers reset(), which the base class implements through set_state.
rows = np.asarray(env_indices, dtype=np.intp)
if rows.ndim == 1 and rows.size:
self._playback_time_view()[rows] = 0.0
return result
def _validate_playback_joint_order(self) -> None:
"""Fail closed when the playback model's MJCF joint order differs from the layout.
Physics-state snapshots serve the worker-maintained public
``qpos``/``qvel`` shared-memory rows, whose columns are only meaningful
to a MuJoCo playback shell while the composed scene file returned by
``get_playback_model`` declares the same joints at the same addresses.
The stdlib re-parse mirrors MuJoCo's depth-first generalized-state
ordering (it is verified against mujoco itself in the Motrix suite), so
any drift between the composed file and the compiled scene layout fails
construction here, before a worker is spawned. Fixed-variant sources
share the canonical layout by the variant handshake, so validating the
canonical composed file covers them.
"""
from unisim.backend.mjcf_layout import extract_mjcf_joint_layout
assert self._entity_scene is not None
layout = self._entity_scene.layout
entries = extract_mjcf_joint_layout(self._entity_scene.owner.model_file)
by_qpos_address = {entry.qpos_address: entry for entry in entries}
expected: list[tuple[tuple[int, ...], tuple[int, ...], str, str]] = []
for entity in layout.entities:
if entity.root_mode == "floating":
expected.append(
(
entity.root_qpos_indices,
entity.root_qvel_indices,
"free",
f"{entity.name}/{entity.root_body}",
)
)
for joint in entity.joints:
expected.append(
(
joint.qpos_indices,
joint.qvel_indices,
joint.kind,
f"{entity.name}/{joint.name}",
)
)
if len(entries) != len(expected):
raise RuntimeError(
f"IsaacSim playback model joint inventory ({len(entries)}) differs from "
f"the scene layout ({len(expected)})"
)
for qpos_indices, qvel_indices, kind, label in expected:
entry = by_qpos_address.get(qpos_indices[0])
if entry is None:
raise RuntimeError(
f"IsaacSim playback model is missing the scene-layout joint {label!r} "
f"at qpos address {qpos_indices[0]}"
)
if (
entry.kind != kind
or entry.qvel_address != qvel_indices[0]
or entry.num_dof_pos != len(qpos_indices)
or entry.num_dof_vel != len(qvel_indices)
or (entry.body_name if kind == "free" else entry.name) != label
):
raise RuntimeError(
"IsaacSim playback model joint order differs from the scene layout at "
f"qpos address {qpos_indices[0]}: {label!r}"
)
def _resolve_model_file_playback_layout(self) -> PhysicsStateLayout | None:
"""Return the model-file snapshot layout, or ``None`` when the wire cannot replay.
The model-file whole-model wire always carries a synthetic 7/6 floating
root followed by the worker's scalar dofs in MJCF document order (the
INIT handshake already validates that order against the worker). A
detached MuJoCo playback shell can consume snapshot rows only while
the source MJCF declares exactly that generalized layout: one named
freejoint body at address zero, then the metadata joints in document
order, all scalar hinge/slide. This is the same hard joint-order
validation as the Motrix non-portable profile, but it reports
``None`` instead of raising: scenes that diverge (fixed base, ball
joints, multiple or reordered roots) still train on the model-file wire,
so only the playback capability fails closed.
"""
from unisim.backend.mjcf_layout import extract_mjcf_joint_layout
metadata = self._get_scene_metadata()
try:
entries = extract_mjcf_joint_layout(str(self._scene.model_file))
except (ValueError, NotImplementedError):
return None
joint_names = tuple(metadata.joint_names)
if len(entries) != 1 + len(joint_names):
return None
root = entries[0]
if (
root.kind != "free"
or root.qpos_address != 0
or root.qvel_address != 0
or metadata.freejoint_body_name is None
or root.body_name != metadata.freejoint_body_name
):
return None
for index, (entry, name) in enumerate(zip(entries[1:], joint_names, strict=True)):
if (
entry.kind not in ("hinge", "slide")
or entry.name != name
or entry.qpos_address != 7 + index
or entry.qvel_address != 6 + index
or entry.num_dof_pos != 1
or entry.num_dof_vel != 1
):
return None
return PhysicsStateLayout(nq=7 + len(joint_names), nv=6 + len(joint_names))
[docs]
def resolve_play_render_plan(
self,
*,
play_render_mode: str | None,
play_steps: int | None,
output_video: str | os.PathLike[str] | None,
) -> PlayRenderPlan:
requested_mode = normalize_play_render_mode(play_render_mode)
if requested_mode == "none":
return PlayRenderPlan(
mode="none",
headless=True,
record_video=False,
num_steps=None,
output_video=None,
)
# Kit's experience is immutable after AppLauncher starts. Resolve
# ``auto`` from the same cold-start decision and fail before playback
# if a direct caller created a no-rendering or differently configured
# worker. The eval adapters inject matching intent before env creation.
startup_mode = self._resolved_render_mode
if startup_mode is None:
startup_mode = self._resolve_render_mode()
mode = startup_mode if requested_mode == "auto" else requested_mode
if startup_mode == "none" or startup_mode != mode:
raise IsaacSimRenderError(
f"IsaacSim worker started in render_mode={startup_mode!r}, but playback requested "
f"{requested_mode!r}; select the matching training.play_render_mode "
"before env creation."
)
if mode == "interactive":
return PlayRenderPlan(
mode=mode,
headless=False,
record_video=False,
num_steps=None,
output_video=None,
)
if isinstance(play_steps, bool) or play_steps is None or int(play_steps) <= 0:
raise ValueError(
"isaacsim record playback requires a positive finite training.play_steps value."
)
if output_video is None:
raise ValueError("isaacsim record playback requires an output video path.")
return PlayRenderPlan(
mode="record",
headless=True,
record_video=True,
num_steps=int(play_steps),
output_video=output_video,
)
[docs]
def bind_renderer(
self,
spacing: float = 1.0,
*,
offset_mode: str = "grid",
headless: bool = False,
capture: bool = False,
width: int = 1280,
height: int = 720,
camera_kwargs: CameraCfg | Mapping[str, Any] | None = None,
) -> None:
mode = self._resolved_render_mode
if mode is None:
mode = self._resolve_render_mode()
requested = "record" if (headless or capture) else "interactive"
if mode != requested:
raise IsaacSimRenderError(
f"IsaacSim worker started in render_mode={mode!r}, but renderer requested "
f"{requested!r}; select the matching training.play_render_mode before env creation."
)
if (
isinstance(width, bool)
or not isinstance(width, (int, np.integer))
or isinstance(height, bool)
or not isinstance(height, (int, np.integer))
or int(width) <= 0
or int(height) <= 0
):
raise IsaacSimRenderError(
f"IsaacSim render dimensions must be positive integers; got {width!r}x{height!r}."
)
if int(width) != self._render_width or int(height) != self._render_height:
raise IsaacSimRenderError(
"IsaacSim render dimensions are fixed at worker startup: "
f"configured {self._render_width}x{self._render_height}, "
f"requested {width}x{height}."
)
# ``super`` owns the protocol/lifecycle and is intentionally called
# only after the mode/dimension checks above.
super().bind_renderer(
spacing=spacing,
offset_mode=offset_mode,
headless=headless,
capture=capture,
width=width,
height=height,
camera_kwargs=camera_kwargs,
)
[docs]
def capture_video_frame(self) -> np.ndarray:
frame = np.asarray(super().capture_video_frame())
expected = (self._render_height, self._render_width, 3)
if frame.dtype != np.uint8 or frame.shape != expected:
raise IsaacSimRenderError(
"IsaacSim camera returned an invalid RGB frame: "
f"shape={frame.shape}, dtype={frame.dtype}, expected shape={expected}, dtype=uint8"
)
if frame.size == 0 or int(np.ptp(frame)) == 0:
raise IsaacSimRenderError(
"IsaacSim camera returned an empty/uniform RGB frame; refusing to write a "
"placeholder video. Check camera pose, lighting, and RTX camera support."
)
return np.ascontiguousarray(frame)
# The model metadata shape is backend-neutral; retaining the alias avoids a
# second, structurally identical dataclass while the error class above remains
# intentionally distinct.
__all__ = [
"IsaacSimBackend",
"IsaacSimModelInfo",
"IsaacSimRenderError",
"IsaacSimWorkerError",
]