"""Mapped IsaacLab scene execution, loaded by path in the external SDK worker.
Asset parsing, USD edits and native-name discovery are cold-path operations.
Runtime operations use frozen entity/view maps, never actor creation order.
"""
from __future__ import annotations
import json
import os
import platform
import re
import tempfile
import time
from hashlib import sha256
from importlib.metadata import PackageNotFoundError, version
from pathlib import Path, PurePosixPath
from typing import Any, cast
import numpy as np
from unisim.backend.isaacsim.physx_solver import (
PHYSX_SOLVER_AUTHORED_FIELDS,
PhysxSolverConfig,
apply_max_depenetration_velocity,
build_isaaclab_physx_cfg,
read_engine_solver_values,
)
from unisim.backend.isaacsim.raw_usd_cache import (
RAW_USD_ARTIFACT_STAGE,
ROLE_USD_ARTIFACT_STAGE,
ROLE_USD_CACHE_SCHEMA_VERSION,
RawUSDArtifactRequest,
RawUSDCache,
RoleUSDCache,
file_sha256,
raw_artifact_fingerprint,
)
from unisim.backend.isaacsim.tensor_ipc import (
WorkerCudaIpcArena,
require_cuda_tensor,
require_cuda_tensor_of_dtype,
)
from unisim.backend.subprocess_ipc.scene_materialization import (
body_sphere_radii_close,
validate_body_sphere_radii,
)
from unisim.scene_compiler import (
SceneContentIdentity,
derive_scene_artifact_identity,
)
def _numpy(value: Any) -> np.ndarray:
if isinstance(value, np.ndarray):
return value.astype(np.float32, copy=False)
return value.detach().cpu().numpy().astype(np.float32, copy=False)
def _rotate(q: np.ndarray, v: np.ndarray, inverse: bool = False) -> np.ndarray:
"""Rotate vectors with the canonical wxyz formula, optionally inversely."""
q = np.asarray(q, dtype=np.float64).copy()
if inverse:
q[..., 1:] *= -1
v = np.asarray(v, dtype=np.float64)
t = 2 * np.cross(q[..., 1:], v)
return (v + q[..., :1] * t + np.cross(q[..., 1:], t)).astype(np.float32)
def _entity_prim_component(name: str) -> str:
"""Injective USD identifier; public names never become native path syntax."""
return "entity_" + name.encode("utf-8").hex()
def _parse_env_number(path: str) -> int:
"""Parse the ``env_<i>`` number from a generated environment path."""
prefix = "/World/envs/env_"
if not path.startswith(prefix):
raise RuntimeError("native view contains an unowned or ambiguous instance")
digits = path[len(prefix) :].split("/", 1)[0]
if not digits.isdigit():
raise RuntimeError("native view contains an unowned or ambiguous instance")
return int(digits)
def _environment_positions(entity_paths: list[str]) -> dict[int, int]:
"""Map env_<i> numbers to entity path positions, failing closed on duplicates."""
positions = {}
for position, root in enumerate(entity_paths):
number = _parse_env_number(root)
if number in positions:
raise RuntimeError("native view contains an unowned or ambiguous instance")
positions[number] = position
return positions
def _environment_index(path: str, entity_paths: list[str], positions: dict[int, int]) -> int:
"""Resolve one path to its environment position via the env_<i> structure.
The resolved entity root prefix is still verified exactly, and anything
unowned or malformed fails closed; semantics match the former quadratic
prefix scan bit for bit.
"""
position = positions.get(_parse_env_number(path))
if position is None:
raise RuntimeError("native view contains an unowned or ambiguous instance")
root = entity_paths[position]
if path != root and not path.startswith(root + "/"):
raise RuntimeError("native view contains an unowned or ambiguous instance")
return position
def _assignment_representatives(observed: list[int]) -> dict[int, int]:
"""Return the first environment index per assigned variant (audit representatives)."""
representatives: dict[int, int] = {}
for env_index, variant in enumerate(observed):
representatives.setdefault(variant, env_index)
return representatives
def _native_environment_order(native_paths: list[str], entity_paths: list[str]) -> np.ndarray:
"""Resolve view rows against exact entity subtrees, never string prefixes alone."""
positions = _environment_positions(entity_paths)
native_envs = [_environment_index(path, entity_paths, positions) for path in native_paths]
if sorted(native_envs) != list(range(len(entity_paths))):
raise RuntimeError("native view needs exactly one instance per environment")
return np.asarray(native_envs, dtype=np.int64)
def _validated_assignment(entry: dict[str, Any], count: int) -> np.ndarray:
"""Validate the immutable per-environment variant references before Kit starts."""
assignment = np.asarray(entry["assignment"])
if (
assignment.shape != (count,)
or assignment.dtype.kind not in "iu"
or bool(np.issubdtype(assignment.dtype, np.bool_))
):
raise ValueError(
"entity assignment must be an integer array with one value per environment"
)
values = assignment.astype(np.int64, copy=False)
source_count = len(entry["sources"])
if values.min(initial=0) < 0 or values.max(initial=-1) >= source_count:
raise ValueError("entity assignment contains an invalid variant index")
return values
def _required_package_version(name: str) -> str:
try:
return str(version(name))
except PackageNotFoundError as exc:
raise RuntimeError(f"IsaacSim worker is missing required package {name!r}") from exc
def _raw_usd_runtime_versions() -> dict[str, str]:
"""Record every importer/runtime input that can change raw USD semantics."""
import omni.kit.app
extension_name = "isaacsim.asset.importer.mjcf"
manager = omni.kit.app.get_app().get_extension_manager()
extension_id = manager.get_enabled_extension_id(extension_name)
if not isinstance(extension_id, str) or not extension_id.startswith(extension_name + "-"):
raise RuntimeError(f"IsaacSim extension {extension_name!r} is unavailable")
extension_version = extension_id[len(extension_name) + 1 :]
if not isinstance(extension_version, str) or not extension_version:
raise RuntimeError(f"IsaacSim extension {extension_name!r} has no readable version")
return {
"isaacsim": _required_package_version("isaacsim"),
"isaaclab": _required_package_version("isaaclab"),
"isaacsim.asset.importer.mjcf": extension_version,
"python": platform.python_version(),
}
def _raw_usd_request(
content_identity: SceneContentIdentity,
source: str,
entity: Any,
variant: int,
runtime_versions: dict[str, str],
*,
self_collision: bool,
) -> RawUSDArtifactRequest:
source_digest = file_sha256(Path(source))
parameters = {
"entity": entity.name,
"variant": variant,
"expanded_source_sha256": source_digest,
"converter": "MjcfConverter",
"importer": {
"fix_base": entity.kind == "articulation" and entity.root_mode == "fixed",
"import_sites": False,
"import_inertia_tensor": True,
"link_density": 0.0,
"make_instanceable": False,
# Cache identity must track the actual converter input exactly.
"self_collision": self_collision,
"force_usd_conversion": True,
"usd_file": "artifact.usd",
},
}
identity = derive_scene_artifact_identity(
content_identity,
RAW_USD_ARTIFACT_STAGE,
{"parameters": parameters, "runtime_versions": runtime_versions},
)
return RawUSDArtifactRequest(identity, source_digest, parameters, runtime_versions)
def _role_usd_request(
raw_record: Any,
entity: Any,
entry: dict[str, Any],
variant: int,
*,
require_bodies: bool,
contact_offset: float | None,
rest_offset: float | None,
share_friction_materials: bool = False,
) -> RawUSDArtifactRequest:
"""Derive one immutable role artifact from its raw identity and bake inputs."""
parameters: dict[str, Any] = {
"entity": entity.name,
"source_entity": entry["mirror_of"] or entity.name,
"variant": variant,
"kind": entity.kind,
"root_mode": entity.root_mode,
"collision_enabled": bool(entry["collision_enabled"]),
"mirror": entry["mirror_of"] is not None,
"mirror_of": entry["mirror_of"],
"raw_identity": raw_record.identity,
"raw_artifact_sha256": raw_artifact_fingerprint(raw_record),
"runtime_versions": dict(raw_record.runtime_versions),
"bake": {
"schema_version": ROLE_USD_CACHE_SCHEMA_VERSION,
"variant_metadata": True,
"remove_joints": entity.kind == "rigid",
"articulation_root": "root-prim" if entity.root_mode == "fixed" else "imported",
"disable_converter_drives": True,
# Opt-in only: equal initial friction values share one PhysX
# material. Reset writes must preserve that invariant fail-closed.
"share_friction_materials": share_friction_materials,
# Resolved per-entity request; cache identity tracks it exactly.
"disable_gravity": bool(entry["gravity_disabled"]),
"require_native_body_paths": require_bodies,
"contact_reporter_api": require_bodies,
# Entity-level solver offsets are baked into the role artifact;
# the cache identity tracks them exactly.
"contact_offset": contact_offset,
"rest_offset": rest_offset,
},
}
placeholders = None
variants = entry.get("variants")
if isinstance(variants, list) and 0 <= variant < len(variants):
placeholders = variants[variant].get("geom_placeholders")
if placeholders and any(placeholders):
# Placeholder prims are baked into the role artifact; the padded slot
# set depends on the compiled public layout, which the raw identity
# does not cover, so the cache identity tracks it exactly.
parameters["bake"]["placeholder_geoms"] = [
index for index, flag in enumerate(placeholders) if flag
]
identity = sha256(
_canonical_role_json(
{
"schema_version": ROLE_USD_CACHE_SCHEMA_VERSION,
"stage": ROLE_USD_ARTIFACT_STAGE,
"raw_identity": raw_record.identity,
"parameters": parameters,
}
).encode("utf-8")
).hexdigest()
return RawUSDArtifactRequest(
identity, raw_record.source_digest, parameters, raw_record.runtime_versions
)
def _canonical_role_json(value: dict[str, Any]) -> str:
return json.dumps(
value, sort_keys=True, separators=(",", ":"), allow_nan=False, ensure_ascii=False
)
def _prototype_spawn_paths(component: str, variant_count: int) -> list[str]:
"""Return one off-stage prototype path per unique converted variant."""
return [
f"/World/unisim_prototypes/{component}/{component}_{index}"
for index in range(variant_count)
]
def _assignment_groups(
assignment: np.ndarray, source_count: int, env_paths: list[str]
) -> tuple[tuple[str, ...], ...]:
"""Group exact environment destinations by prototype without expanding K sources."""
if assignment.shape != (len(env_paths),) or assignment.min(initial=0) < 0:
raise ValueError("assignment and environment path counts or values are invalid")
if assignment.max(initial=-1) >= source_count:
raise ValueError("assignment contains an unknown prototype index")
return tuple(
tuple(env_paths[int(row)] for row in np.flatnonzero(assignment == variant))
for variant in range(source_count)
)
def _copy_prims_from_source(
stage: Any, source_prim_path: str, destinations: tuple[str, ...]
) -> None:
"""Copy one prim tree onto destination paths in the stage root layer.
Replicates exactly what ``Cloner.clone`` does per destination with
``copy_from_source=True`` and no explicit positions: the source xform ops
are collapsed to translate/orient/scale, each destination receives a full
``Sdf.CopySpec`` of the source spec, and the destination translate/orient/
scale specs are set to the source's local-to-parent transform. IsaacLab's
``MultiUsdFileCfg`` spawner batches copies the same way. Unlike
``Cloner.clone`` this helper performs no change-listener toggle and opens
no ``Sdf.ChangeBlock``; the caller wraps any number of copies in one
toggle and one block instead of paying both per prototype.
"""
import carb.settings
from pxr import Gf, Sdf, Usd, UsdGeom, Vt
source_prim = stage.GetPrimAtPath(source_prim_path)
if not source_prim:
raise RuntimeError(f"IsaacSim prototype is missing: {source_prim_path}")
properties = source_prim.GetPropertyNames()
xformable = UsdGeom.Xformable(source_prim)
parent_to_world = xformable.ComputeParentToWorldTransform(Usd.TimeCode.Default())
local_to_world = xformable.ComputeLocalToWorldTransform(Usd.TimeCode.Default())
local_to_parent = Gf.Transform()
local_to_parent.SetMatrix(
Gf.Matrix4d(np.matmul(local_to_world, np.linalg.inv(parent_to_world)).tolist())
)
current_translation = local_to_parent.GetTranslation()
current_orientation = local_to_parent.GetRotation().GetQuat()
current_scale = Gf.Vec3d(1, 1, 1)
if "xformOp:scale" in properties:
current_scale = Gf.Vec3d(source_prim.GetAttribute("xformOp:scale").Get())
# Collapse the source xform ops to translate/orient/scale, as Cloner.clone
# does before copying, so destinations carry one canonical op order.
properties_to_remove = [
"xformOp:rotateX",
"xformOp:rotateXZY",
"xformOp:rotateY",
"xformOp:rotateYXZ",
"xformOp:rotateYZX",
"xformOp:rotateZ",
"xformOp:rotateZYX",
"xformOp:rotateZXY",
"xformOp:rotateXYZ",
"xformOp:transform",
"xformOp:scale",
]
xformable.ClearXformOpOrder()
for prop_name in properties:
if prop_name in properties_to_remove:
source_prim.RemoveProperty(prop_name)
properties = source_prim.GetPropertyNames()
if "xformOp:translate" not in properties:
xform_op_translate = xformable.AddXformOp(
UsdGeom.XformOp.TypeTranslate, UsdGeom.XformOp.PrecisionDouble, ""
)
else:
xform_op_translate = UsdGeom.XformOp(source_prim.GetAttribute("xformOp:translate"))
xform_op_translate.Set(current_translation)
if "xformOp:orient" not in properties:
xform_op_rot = xformable.AddXformOp(
UsdGeom.XformOp.TypeOrient, UsdGeom.XformOp.PrecisionDouble, ""
)
else:
xform_op_rot = UsdGeom.XformOp(source_prim.GetAttribute("xformOp:orient"))
if xform_op_rot.GetPrecision() == UsdGeom.XformOp.PrecisionFloat:
current_orientation = Gf.Quatf(current_orientation)
else:
current_orientation = Gf.Quatd(current_orientation)
xform_op_rot.Set(current_orientation)
if "xformOp:scale" not in properties:
xform_op_scale = xformable.AddXformOp(
UsdGeom.XformOp.TypeScale, UsdGeom.XformOp.PrecisionDouble, ""
)
else:
xform_op_scale = UsdGeom.XformOp(source_prim.GetAttribute("xformOp:scale"))
xform_op_scale.Set(current_scale)
xformable.SetXformOpOrder([xform_op_translate, xform_op_rot, xform_op_scale])
root_layer = stage.GetRootLayer()
for prim_path in destinations:
if prim_path == source_prim_path:
continue
destination_spec = Sdf.CreatePrimInLayer(root_layer, prim_path)
Sdf.CopySpec(root_layer, Sdf.Path(source_prim_path), root_layer, Sdf.Path(prim_path))
# CopySpec already carried the normalized source's xformOpOrder and
# op defaults, so rewriting translate/orient/scale/order below is
# redundant; it is kept intentionally to stay byte-aligned with
# Cloner.clone(copy_from_source=True) destination output.
translate_spec = destination_spec.GetAttributeAtPath(prim_path + ".xformOp:translate")
if translate_spec is None:
translate_spec = Sdf.AttributeSpec(
destination_spec, "xformOp:translate", Sdf.ValueTypeNames.Double3
)
translate_spec.default = current_translation
orient_spec = destination_spec.GetAttributeAtPath(prim_path + ".xformOp:orient")
default_precision = carb.settings.get_settings().get_as_string(
"app/primCreation/DefaultXformOpPrecision"
)
if orient_spec is None:
if len(default_precision) > 0 and default_precision == "Float":
orient_spec = Sdf.AttributeSpec(
destination_spec, "xformOp:orient", Sdf.ValueTypeNames.Quatf
)
orient_spec.default = Gf.Quatf(current_orientation)
else:
orient_spec = Sdf.AttributeSpec(
destination_spec, "xformOp:orient", Sdf.ValueTypeNames.Quatd
)
orient_spec.default = Gf.Quatd(current_orientation)
elif orient_spec.default is not None and type(orient_spec.default) is Gf.Quatf:
orient_spec.default = Gf.Quatf(current_orientation)
else:
orient_spec.default = Gf.Quatd(current_orientation)
scale_spec = destination_spec.GetAttributeAtPath(prim_path + ".xformOp:scale")
if scale_spec is None:
scale_spec = Sdf.AttributeSpec(
destination_spec, "xformOp:scale", Sdf.ValueTypeNames.Double3
)
scale_spec.default = current_scale
op_order_spec = destination_spec.GetAttributeAtPath(prim_path + ".xformOpOrder")
if op_order_spec is None:
op_order_spec = Sdf.AttributeSpec(
destination_spec, UsdGeom.Tokens.xformOpOrder, Sdf.ValueTypeNames.TokenArray
)
op_order_spec.default = Vt.TokenArray(
["xformOp:translate", "xformOp:orient", "xformOp:scale"]
)
def _current_rss_mb() -> float | None:
"""Read the worker's current resident set from procfs (Linux-only worker)."""
try:
with open("/proc/self/status", encoding="utf-8") as handle:
for line in handle:
if line.startswith("VmRSS:"):
return float(line.split()[1]) / 1024.0
except OSError:
pass
return None
class _InitTelemetry:
"""Cold-path INIT phase timing with resource snapshots.
``mark()`` closes the span since the previous mark under a phase name
(accumulating, so per-variant loop iterations aggregate into one entry),
``span()`` records an explicitly-timed region such as a whole entity
loop, and ``count_prims()`` snapshots the stage size at structural
boundaries. All of it is INIT-only; step/reset hot paths never touch
this object.
"""
def __init__(self) -> None:
self._t0 = self._mark = time.perf_counter()
self._elapsed_ms: dict[str, float] = {}
self._rss_mb: dict[str, float | None] = {}
self._order: list[str] = []
self.stage_prims: dict[str, int] = {}
self.overhead_ms = 0.0
self.vram_used_mb: float | None = None
def _record(self, name: str, elapsed_ms: float) -> None:
if name not in self._elapsed_ms:
self._order.append(name)
self._elapsed_ms[name] = 0.0
self._elapsed_ms[name] += elapsed_ms
self._rss_mb[name] = _current_rss_mb()
def mark(self, name: str) -> None:
now = time.perf_counter()
self._record(name, (now - self._mark) * 1e3)
self._mark = now
def span(self, name: str, started: float) -> None:
now = time.perf_counter()
self._record(name, (now - started) * 1e3)
self._mark = now
def count_prims(self, stage: Any, key: str) -> None:
started = time.perf_counter()
self.stage_prims[key] = sum(1 for _ in stage.Traverse())
# A full stage traversal scales with the stage itself; account for it
# explicitly instead of billing the following phase.
now = time.perf_counter()
self.overhead_ms += (now - started) * 1e3
self._mark = now
def to_wire(self) -> dict[str, Any]:
import resource
return {
"schema_version": 1,
"init_total_ms": (time.perf_counter() - self._t0) * 1e3,
"phases": [
{
"name": name,
"elapsed_ms": self._elapsed_ms[name],
"rss_mb": self._rss_mb[name],
}
for name in self._order
],
"stage_prims": dict(self.stage_prims),
"peak_rss_mb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0,
"vram_used_mb": self.vram_used_mb,
"telemetry_overhead_ms": self.overhead_ms,
}
def _validated_entry_velocity_limits(entry: dict[str, Any], entity: Any) -> None:
"""Fail closed on a malformed or inconsistent per-joint velocity-limit table.
MJCF sources cannot express joint velocity limits, so the host declaration
is the only carrier; an undeclared entity must not carry record values and
a declared entity must name every compiled joint exactly.
"""
declared = entry.get("joint_velocity_limits")
joint_names = [joint.name for joint in entity.joints]
if declared is None:
for record in entry["variants"]:
if "dof_velocity_limit" in record:
raise ValueError(
"variant dof_velocity_limit requires an entity joint_velocity_limits "
"declaration"
)
return
if not isinstance(declared, dict) or not declared:
raise ValueError("entity joint_velocity_limits must be a non-empty table or None")
if any(not isinstance(name, str) or not name for name in declared):
raise ValueError("entity joint_velocity_limits keys must be non-empty joint names")
if set(declared) != set(joint_names):
raise ValueError("entity joint_velocity_limits must name every joint exactly")
expected = [declared[name] for name in joint_names]
if any(
isinstance(value, (bool, np.bool_))
or not isinstance(value, (int, float, np.integer, np.floating))
or not np.isfinite(float(value))
or float(value) <= 0.0
for value in expected
):
raise ValueError("entity joint_velocity_limits values must be positive and finite")
for record in entry["variants"]:
raw = record.get("dof_velocity_limit")
if not isinstance(raw, list):
raise ValueError("invalid variant dof_velocity_limit")
values = np.asarray(raw, dtype=np.float64)
if values.shape != (len(joint_names),) or not np.isfinite(values).all():
raise ValueError("invalid variant dof_velocity_limit")
if np.any(values <= 0.0):
raise ValueError("variant dof_velocity_limit values must be positive")
if not np.allclose(values, np.asarray(expected, dtype=np.float64), rtol=0.0, atol=0.0):
raise ValueError("variant dof_velocity_limit differs from the entity declaration")
[docs]
def validate_scene_payload(protocol: Any, payload: dict[str, Any]) -> Any:
"""Reject unsupported combinations before launching Kit or converting assets."""
layout = protocol.load_scene_layout(payload["scene_layout"])
SceneContentIdentity.from_dict(payload["scene_content_identity"])
count = payload["num_envs"]
contact_force_sensors = _validate_contact_force_sensors(payload, layout)
protocol.scene_slot_shapes(count, layout, len(contact_force_sensors))
entries = payload["scene_entities"]
if [entry["name"] for entry in entries] != [entity.name for entity in layout.entities]:
raise ValueError("scene entity order differs from the frozen layout")
for entity, entry in zip(layout.entities, entries):
if entry["asset_format"] != "mjcf":
raise NotImplementedError("IsaacSim mapped scene currently requires MJCF sources")
if entity.kind != entry["kind"] or entity.root_mode != entry["root_mode"]:
raise ValueError("entity declaration differs from compiled layout")
self_collision = entry.get("self_collision")
if not isinstance(self_collision, bool):
raise TypeError("entity self_collision must be bool")
if self_collision:
if entity.kind != "articulation":
raise NotImplementedError("IsaacSim self-collision requires an articulation")
if entry["mirror_of"] is not None or not entry["collision_enabled"]:
raise ValueError("self_collision requires a collision-enabled physical entity")
# The host resolves gravity_disabled=None to this backend's implicit
# default before INIT; an unset value here means the request never
# passed through host resolution, so fail closed.
if not isinstance(entry.get("gravity_disabled"), bool):
raise TypeError("entity gravity_disabled must be bool")
_validated_entry_velocity_limits(entry, entity)
sources = entry["sources"]
if not sources or len(sources) != len(entry["variants"]):
raise ValueError("entity source and variant record counts differ")
_validated_assignment(entry, count)
if entity.kind == "rigid" and len(entity.body_names) != 1:
raise NotImplementedError("IsaacSim rigid entity requires one physical body")
if entity.root_mode == "kinematic" and entity.kind != "rigid":
raise NotImplementedError("IsaacSim kinematic articulation is unsupported")
if any(joint.kind == "ball" for joint in entity.joints):
raise NotImplementedError("IsaacSim mapped scene supports scalar joints only")
if len(set(entity.actuator_joint_names)) != len(entity.actuator_joint_names):
raise NotImplementedError("IsaacSim requires one actuator per controlled joint")
for record in entry["variants"]:
if record["joint_names"] != [joint.name for joint in entity.joints]:
raise ValueError("variant joint names differ from compiled layout")
if record["body_names"] != list(entity.body_names):
raise ValueError("variant body names differ from compiled layout")
validate_body_sphere_radii(record["body_sphere_radii"], len(entity.body_names))
# uniform_public_layout permits optional mesh-geom slots to be
# absent from a variant (e.g. eraser heads); the records are
# padded to the full public layout after this validation loop.
# Here, fail closed unless the rows form an order-preserving
# subset of the compiled layout.
public_geoms = list(
zip(
(geom.name for geom in entity.geoms),
(geom.body_name for geom in entity.geoms),
)
)
actual_geoms = list(zip(record["geom_names"], record["geom_body_names"]))
cursor = 0
for pair in actual_geoms:
while cursor < len(public_geoms) and public_geoms[cursor] != pair:
cursor += 1
if cursor == len(public_geoms):
raise ValueError("variant geom names differ from compiled layout")
cursor += 1
for field in ("geom_contype", "geom_conaffinity"):
values = record[field]
if (
not isinstance(values, list)
or len(values) != len(actual_geoms)
or any(
isinstance(value, (bool, np.bool_))
or not isinstance(value, (int, np.integer))
for value in values
)
):
raise ValueError("invalid variant " + field)
_validated_variant_geom_types(record, len(actual_geoms))
_validated_variant_source_geometry(record, len(actual_geoms))
friction = np.asarray(record["geom_friction"], dtype=np.float32)
if (
friction.shape != (len(actual_geoms), 3)
or not np.isfinite(friction).all()
or np.any(friction < 0.0)
):
raise ValueError("invalid variant geom_friction")
if record["actuator_joint_names"] != list(entity.actuator_joint_names):
raise ValueError("variant actuator targets differ from compiled layout")
for field in (
"dof_stiffness",
"dof_damping",
"dof_effort",
"dof_armature",
"dof_friction",
"dof_lower",
"dof_upper",
):
values = np.asarray(record[field])
valid = (
~np.isnan(values)
if field in ("dof_lower", "dof_upper")
else np.isfinite(values)
)
if values.shape != (len(entity.joints),) or not valid.all():
raise ValueError("invalid variant " + field)
for index, joint in enumerate(entity.joints):
if joint.name not in entity.actuator_joint_names and record["dof_stiffness"][index]:
raise NotImplementedError("passive joint stiffness requires explicit semantics")
# Pad slot-omitting variant records to the full public layout so every
# environment realizes one homogeneous native geometry layout.
entry["variants"] = [
_normalized_variant_record(entity, record) for record in entry["variants"]
]
# One view shares one sim-baked effort limit table, and armature,
# joint friction and velocity limits start from a single
# ImplicitActuatorCfg. Drive stiffness/damping may differ across
# variants: initialization rewrites every environment's assigned
# variant gains through the per-row PhysX view setters, and reset DR
# owns later per-env drive updates.
shared_fields = ["dof_effort", "dof_armature", "dof_friction"]
if "dof_velocity_limit" in entry["variants"][0]:
shared_fields.append("dof_velocity_limit")
for record in entry["variants"][1:]:
for field in shared_fields:
if record[field] != entry["variants"][0][field]:
raise NotImplementedError(
"IsaacSim entity variants require identical effort, armature, "
"joint friction and velocity limits; drive stiffness/damping "
"may differ per variant"
)
for field, shape in (
("initial_qpos", (count, layout.nq)),
("initial_qvel", (count, layout.nv)),
("initial_roots", (count, len(layout.entities), 13)),
):
values = np.asarray(payload[field], dtype=np.float32)
if values.shape != shape or not np.isfinite(values).all():
raise ValueError("invalid " + field)
if "initial_ctrl" in payload:
values = np.asarray(payload["initial_ctrl"], dtype=np.float32)
if values.shape != (count, layout.nu) or not np.isfinite(values).all():
raise ValueError("invalid initial_ctrl")
return layout
def _self_collision_entity_names(payload: dict[str, Any]) -> set[str]:
"""Entities whose PhysX contacts include self-contacts by request."""
return {
entry["name"]
for entry in payload.get("scene_entities", [])
if isinstance(entry, dict)
and isinstance(entry.get("name"), str)
and entry.get("self_collision") is True
}
def _validate_contact_force_sensors(payload: dict[str, Any], layout: Any) -> list[dict[str, str]]:
records = payload.get("contact_force_sensors", [])
if not isinstance(records, list):
raise ValueError("contact_force_sensors must be a list")
entities = {entity.name: entity for entity in layout.entities}
self_colliding = _self_collision_entity_names(payload)
names: list[str] = []
body_pairs: dict[tuple[str, str, str, str], str] = {}
for record in records:
if not isinstance(record, dict) or set(record) != {
"name",
"source_entity",
"source_body",
"target_entity",
"target_body",
}:
raise ValueError("malformed contact force sensor declaration")
if not all(isinstance(record[key], str) and record[key] for key in record):
raise ValueError("contact force sensor fields must be non-empty strings")
for role in ("source", "target"):
entity = entities.get(record[f"{role}_entity"])
body = record[f"{role}_body"]
if entity is None or body not in entity.body_names:
raise ValueError(
f"contact force sensor {record['name']!r} references unknown "
f"{role} entity/body: {record[f'{role}_entity']}/{body}"
)
if (
record["source_entity"] == record["target_entity"]
and record["source_entity"] in self_colliding
):
# The filtered pair reporter cannot isolate self-contacts from the
# requested per-entity self-collision, so the combination fails
# closed instead of reporting forces the declaration did not mean.
raise ValueError(
f"contact force sensor {record['name']!r} measures a same-entity "
"pair on a self-collision entity"
)
if record["name"] in names:
raise ValueError("duplicate contact force sensor name: " + record["name"])
names.append(record["name"])
body_pair = (
record["source_entity"],
record["source_body"],
record["target_entity"],
record["target_body"],
)
previous_name = body_pairs.get(body_pair)
if previous_name is not None:
raise ValueError(
"duplicate ordered rigid-body pair for contact force sensors: "
f"{previous_name!r} and {record['name']!r}"
)
body_pairs[body_pair] = record["name"]
return records
def _validate_body_net_contact_entities(payload: dict[str, Any], layout: Any) -> list[str]:
"""Validate the host's per-body net contact force coverage request."""
names = payload.get("body_net_contact_entities", [])
if not isinstance(names, list) or any(not isinstance(name, str) or not name for name in names):
raise ValueError("body_net_contact_entities must be a list of non-empty strings")
if len(set(names)) != len(names):
raise ValueError("body_net_contact_entities must not contain duplicates")
known = {entity.name for entity in layout.entities}
unknown = [name for name in names if name not in known]
if unknown:
raise ValueError(f"body_net_contact_entities references unknown entities: {unknown}")
self_colliding = sorted(name for name in names if name in _self_collision_entity_names(payload))
if self_colliding:
# PhysX net contact forces of a self-collision entity would include
# self-contacts, contradicting the wildcard declaration's semantics.
raise ValueError(
f"body_net_contact_entities includes self-collision entities: {self_colliding}"
)
return names
def _bake(
usd_path: str,
entity: Any,
entry: dict[str, Any],
variant: int,
body_paths: dict[str, str],
require_bodies: bool = False,
contact_offset: float | None = None,
rest_offset: float | None = None,
share_friction_materials: bool = False,
) -> str:
"""Author declared root/role semantics and immutable source identity on USD."""
from pxr import PhysxSchema, Sdf, Usd, UsdPhysics
stage = Usd.Stage.Open(usd_path)
root = stage.GetDefaultPrim()
if not root or not root.IsValid():
raise RuntimeError("converted entity USD has no default prim")
for prim in Usd.PrimRange(root, Usd.TraverseInstanceProxies()):
if prim.IsInstance():
prim.SetInstanceable(False)
root.CreateAttribute("unisim:variantIndex", Sdf.ValueTypeNames.Int).Set(variant)
articulation_roots = []
rigid_bodies = []
remove_joints = []
root_path = str(root.GetPath())
for prim in Usd.PrimRange(root):
if entity.kind == "rigid" and prim.IsA(UsdPhysics.Joint):
remove_joints.append(str(prim.GetPath()))
if prim.HasAPI(UsdPhysics.ArticulationRootAPI):
if entity.kind == "rigid":
prim.RemoveAPI(UsdPhysics.ArticulationRootAPI)
if prim.HasAPI(PhysxSchema.PhysxArticulationAPI):
prim.RemoveAPI(PhysxSchema.PhysxArticulationAPI)
elif prim.GetName() == entity.root_body:
articulation_roots.append(str(prim.GetPath()))
# The MJCF importer records the converter self_collision input
# as physxArticulation:enabledSelfCollisions on the articulation
# root. Authored <contact><exclude> pairs survive separately as
# UsdPhysics FilteredPairsAPI relationships.
flag = PhysxSchema.PhysxArticulationAPI(prim).GetEnabledSelfCollisionsAttr()
if bool(flag.Get()) != bool(entry["self_collision"]):
raise RuntimeError(
f"entity {entity.name} importer self-collision differs from request"
)
else:
# Importer also tags its synthetic worldBody, outside this entity.
prim.RemoveAPI(UsdPhysics.ArticulationRootAPI)
if prim.HasAPI(UsdPhysics.RigidBodyAPI):
rigid_bodies.append(prim)
body_name = prim.GetName()
prim_path = str(prim.GetPath())
if body_name in body_paths and body_paths[body_name] != prim_path[len(root_path) :]:
raise RuntimeError(f"entity {entity.name} has duplicate rigid body {body_name!r}")
body_paths[body_name] = prim_path[len(root_path) :]
if entity.kind == "rigid":
UsdPhysics.RigidBodyAPI(prim).CreateKinematicEnabledAttr().Set(
entity.root_mode != "floating"
)
PhysxSchema.PhysxRigidBodyAPI.Apply(prim).CreateDisableGravityAttr().Set(
bool(entry["gravity_disabled"])
)
if require_bodies:
# IsaacLab's ContactSensor expects this API on every source rigid
# body. Author it in the immutable role artifact so environment
# copies inherit it and cold cache hits remain valid.
PhysxSchema.PhysxRigidBodyAPI.Apply(prim).CreateSleepThresholdAttr().Set(0.0)
PhysxSchema.PhysxContactReportAPI.Apply(prim).CreateThresholdAttr().Set(0.0)
if prim.HasAPI(UsdPhysics.CollisionAPI):
UsdPhysics.CollisionAPI(prim).CreateCollisionEnabledAttr().Set(
bool(entry["collision_enabled"])
)
# Solver offsets are entity-level configuration identical across
# every environment copy: author them once into the role artifact
# instead of per destination prim on the composed stage.
if contact_offset is not None or rest_offset is not None:
physx_collision = PhysxSchema.PhysxCollisionAPI.Apply(prim)
if contact_offset is not None:
physx_collision.CreateContactOffsetAttr().Set(float(contact_offset))
if rest_offset is not None:
physx_collision.CreateRestOffsetAttr().Set(float(rest_offset))
# Disable converter-authored drives. IsaacLab owns declared control gains.
for axis in ("angular", "linear"):
if prim.HasAPI(UsdPhysics.DriveAPI, axis):
drive = UsdPhysics.DriveAPI(prim, axis)
drive.CreateStiffnessAttr().Set(0.0)
drive.CreateDampingAttr().Set(0.0)
if entity.kind == "rigid" and len(rigid_bodies) != 1:
raise RuntimeError(f"rigid entity {entity.name} has {len(rigid_bodies)} native bodies")
for path in remove_joints:
# Converter prims may be authored in referenced layers. An inactive
# override suppresses them; RemovePrim would merely reveal the reference.
stage.GetPrimAtPath(path).SetActive(False)
missing = [name for name in entity.body_names if name not in body_paths]
if require_bodies and missing:
raise RuntimeError(
f"entity {entity.name} converted rigid-body paths are missing bodies: {missing}"
)
record = entry["variants"][variant]
_author_placeholder_geoms(
stage,
root_path,
body_paths,
record,
contact_offset=contact_offset,
rest_offset=rest_offset,
)
_author_native_geometry(
stage,
root_path,
body_paths,
entity,
record,
share_friction_materials=share_friction_materials,
)
relative = ""
if entity.kind == "articulation":
if len(articulation_roots) != 1:
raise RuntimeError(
f"entity {entity.name} has ambiguous articulation roots: {articulation_roots}"
)
if entity.root_mode == "fixed":
# A world joint attached to an API-bearing rigid link remains a
# floating articulation constrained by a maximal-coordinate joint.
# PhysX recognizes a fixed-base tree when the root API is on the
# encompassing prim instead (IsaacLab's fix_root_link convention).
body_prim = stage.GetPrimAtPath(articulation_roots[0])
body_prim.RemoveAPI(UsdPhysics.ArticulationRootAPI)
if body_prim.HasAPI(PhysxSchema.PhysxArticulationAPI):
body_prim.RemoveAPI(PhysxSchema.PhysxArticulationAPI)
UsdPhysics.ArticulationRootAPI.Apply(root)
# Re-applying the PhysX API on the root prim drops importer-authored
# attributes, so the verified self-collision flag is authored again.
PhysxSchema.PhysxArticulationAPI.Apply(root).CreateEnabledSelfCollisionsAttr().Set(
bool(entry["self_collision"])
)
else:
relative = articulation_roots[0][len(str(root.GetPath())) :]
stage.GetRootLayer().Save()
return relative
def _body_collision_prims(body_prim: Any) -> list[Any]:
"""Return direct collision leaves without crossing a nested rigid body."""
from pxr import UsdPhysics
result = []
def visit(prim: Any, *, root: bool = False) -> None:
if not root and prim.HasAPI(UsdPhysics.RigidBodyAPI):
return
if prim.HasAPI(UsdPhysics.CollisionAPI):
result.append(prim)
return
for child in prim.GetChildren():
visit(child)
visit(body_prim, root=True)
return result
def _native_geometry_columns(
native_body_names: list[str],
body_permutation: np.ndarray,
entity: Any,
collision_mask: np.ndarray | None = None,
) -> np.ndarray:
"""Map public geometry order to PhysX's flattened native shape order.
Visual-only geoms own no native shape (the MJCF importer honors their
zero contact bits); their column is -1 and every consumer must skip it.
``collision_mask=None`` keeps the model-file all-colliding mapping.
"""
permutation = np.asarray(body_permutation, dtype=np.int64)
if permutation.shape != (len(native_body_names),):
raise RuntimeError("native body permutation does not match the rigid-body view")
if collision_mask is None:
collision_mask = np.ones(len(entity.geoms), dtype=bool)
else:
collision_mask = np.asarray(collision_mask, dtype=bool)
if collision_mask.shape != (len(entity.geoms),):
raise RuntimeError("collision mask does not match the public geometry layout")
public_positions = {name: index for index, name in enumerate(entity.body_names)}
native_counts = np.zeros(len(native_body_names), dtype=np.int64)
for geom_index, geom in enumerate(entity.geoms):
public_position = public_positions.get(geom.body_name)
if public_position is None:
raise RuntimeError(
f"entity {entity.name} geometry references unknown body {geom.body_name!r}"
)
if not collision_mask[geom_index]:
continue
native_body = int(permutation[public_position])
native_counts[native_body] += 1
native_starts = np.zeros(len(native_body_names), dtype=np.int64)
if len(native_counts) > 1:
native_starts[1:] = np.cumsum(native_counts[:-1])
local_offsets = np.zeros(len(native_body_names), dtype=np.int64)
columns = np.full(len(entity.geoms), -1, dtype=np.int64)
for geom_index, geom in enumerate(entity.geoms):
if not collision_mask[geom_index]:
continue
native_body = int(permutation[public_positions[geom.body_name]])
columns[geom_index] = native_starts[native_body] + local_offsets[native_body]
local_offsets[native_body] += 1
return columns
def _record_collision_mask(record: dict[str, Any]) -> np.ndarray:
"""Per-geom collision participation from a variant record.
A geom collides when either contact bit is nonzero (MuJoCo semantics);
``contype=0`` and ``conaffinity=0`` together mark a visual-only geom.
The MJCF importer honors those bits and authors no CollisionAPI for
visual-only geoms, so every native-geometry consumer below works on the
colliding subset while keeping the public (full) geom ordering.
"""
contype = record["geom_contype"]
conaffinity = record["geom_conaffinity"]
return np.asarray(
[int(ct) != 0 or int(ca) != 0 for ct, ca in zip(contype, conaffinity)],
dtype=bool,
)
def _record_native_mask(record: dict[str, Any]) -> np.ndarray:
"""Per-slot native shape presence from a (possibly padded) variant record.
Placeholder slots own a collision-disabled placeholder prim, so they count
toward the native geometry layout even though they never collide.
"""
mask = _record_collision_mask(record)
placeholders = record.get("geom_placeholders")
if placeholders is None:
return mask
return mask | np.asarray(placeholders, dtype=bool)
def _validated_variant_geom_types(record: dict[str, Any], expected_count: int) -> list[str]:
"""Validate the required source-geom type column.
``geom_types`` is generated only by the current host materializer and is
deliberately not optional at the worker boundary. It distinguishes
source-authored spheres from USD sphere approximations, so accepting an
older payload would silently change geometry audit semantics.
"""
if "geom_types" not in record:
raise ValueError("variant geom_types are required by the current protocol")
values = record["geom_types"]
if (
not isinstance(values, list)
or len(values) != expected_count
or any(not isinstance(value, str) or value not in ("sphere", "other") for value in values)
):
raise ValueError(
"invalid variant geom_types; expected a list of 'sphere' or 'other' "
f"with length {expected_count}"
)
return values
def _validated_variant_source_geometry(record: dict[str, Any], expected_count: int) -> None:
"""Validate immutable source geometry intent in public geom order.
Source size uses MuJoCo's ``[x, y, z]`` columns and source pose uses
``[x, y, z, qw, qx, qy, qz]``. Keeping the raw columns at the worker
boundary lets the cold native audit compare importer output to source
intent without parsing MJCF on the hot path.
"""
for field in ("geom_source_types", "geom_source_sizes", "geom_source_poses"):
if field not in record:
raise ValueError("variant " + field + " are required by the current protocol")
types = record["geom_source_types"]
if (
not isinstance(types, list)
or len(types) != expected_count
or any(
not isinstance(value, str) or not re.fullmatch(r"[a-z][a-z0-9_]*", value)
for value in types
)
):
raise ValueError("invalid variant geom_source_types")
try:
sizes = np.asarray(record["geom_source_sizes"], dtype=np.float64)
poses = np.asarray(record["geom_source_poses"], dtype=np.float64)
except (TypeError, ValueError) as exc:
raise ValueError("invalid variant geom_source_sizes or geom_source_poses") from exc
if (
sizes.shape != (expected_count, 3)
or poses.shape != (expected_count, 7)
or not np.isfinite(sizes).all()
or not np.isfinite(poses).all()
or np.any(sizes < 0.0)
):
raise ValueError("invalid variant geom_source_sizes or geom_source_poses")
quaternion_norm = np.linalg.norm(poses[:, 3:7], axis=1)
if not np.allclose(quaternion_norm, 1.0, rtol=0.0, atol=1e-5):
raise ValueError("invalid variant geom_source_poses quaternion")
def _finite_matrix(matrix: Any, label: str) -> list[list[float]]:
try:
values = np.asarray(matrix, dtype=np.float64)
except (TypeError, ValueError) as exc:
raise RuntimeError(f"native {label} is malformed") from exc
if (
values.shape != (4, 4)
or not np.isfinite(values).all()
or not np.allclose(values[:, 3], (0.0, 0.0, 0.0, 1.0), rtol=0.0, atol=1e-6)
):
raise RuntimeError(
f"native {label} is malformed: shape={values.shape}, values={values.tolist()!r}"
)
return cast(list[list[float]], values.tolist())
def _range_bounds(bounds: Any, label: str) -> tuple[bool, list[float] | None, list[float] | None]:
if bounds is None:
raise RuntimeError(f"native {label} bounds are missing")
try:
if bounds.IsEmpty():
return True, None, None
minimum = np.asarray(bounds.GetMin(), dtype=np.float64)
maximum = np.asarray(bounds.GetMax(), dtype=np.float64)
except (AttributeError, TypeError, ValueError) as exc:
raise RuntimeError(f"native {label} bounds are malformed") from exc
if (
minimum.shape != (3,)
or maximum.shape != (3,)
or not np.isfinite(minimum).all()
or not np.isfinite(maximum).all()
or np.any(maximum < minimum)
):
raise RuntimeError(f"native {label} bounds are malformed")
return False, minimum.tolist(), maximum.tolist()
def _mesh_summary(mesh: Any) -> dict[str, Any]:
points_attr = mesh.GetPointsAttr()
counts_attr = mesh.GetFaceVertexCountsAttr()
indices_attr = mesh.GetFaceVertexIndicesAttr()
if points_attr is None or counts_attr is None or indices_attr is None:
raise RuntimeError("native collision mesh attributes are missing")
points = points_attr.Get()
counts = counts_attr.Get()
indices = indices_attr.Get()
try:
point_values = np.asarray(points, dtype=np.float64)
count_values = np.asarray(counts, dtype=np.int64)
index_values = np.asarray(indices, dtype=np.int64)
except (TypeError, ValueError) as exc:
raise RuntimeError("native collision mesh attributes are malformed") from exc
if (
point_values.size == 0
or point_values.ndim != 2
or point_values.shape[1] != 3
or count_values.ndim != 1
or index_values.ndim != 1
or not np.isfinite(point_values).all()
or np.any(count_values < 3)
or index_values.size != int(count_values.sum(initial=0))
or np.any(index_values < 0)
or (point_values.shape[0] and np.any(index_values >= point_values.shape[0]))
):
raise RuntimeError("native collision mesh attributes are malformed")
tolerance = 1e-12
# Canonicalize signed zero before the exact unique-vertex count.
point_values[point_values == 0.0] = 0.0
point_extent = point_values.max(axis=0) - point_values.min(axis=0)
unique_count = len(np.unique(point_values, axis=0))
all_coincident = bool(np.all(point_extent <= tolerance))
surface_area = 0.0
index_cursor = 0
for count in count_values:
polygon = index_values[index_cursor : index_cursor + int(count)]
for first, second in zip(polygon[1:-1], polygon[2:], strict=False):
edge_a = point_values[int(first)] - point_values[int(polygon[0])]
edge_b = point_values[int(second)] - point_values[int(polygon[0])]
surface_area += float(np.linalg.norm(np.cross(edge_a, edge_b)) / 2.0)
index_cursor += int(count)
return {
"mesh_vertex_count": int(point_values.shape[0]),
"mesh_unique_vertex_count": unique_count,
"mesh_face_count": int(count_values.shape[0]),
"mesh_face_vertex_index_count": int(index_values.shape[0]),
"mesh_tolerance": tolerance,
"mesh_all_points_coincident": all_coincident,
"mesh_points_min": point_values.min(axis=0).tolist(),
"mesh_points_max": point_values.max(axis=0).tolist(),
"mesh_points_extent": point_extent.tolist(),
"mesh_points_zero_extent": all_coincident,
"mesh_surface_area": surface_area,
"mesh_zero_area": surface_area <= tolerance * tolerance,
}
def _native_collision_geometry_audit(collision: Any, bbox_cache: Any) -> dict[str, Any]:
"""Read one actual USD collision leaf without interpreting source intent."""
from pxr import PhysxSchema, Usd, UsdGeom, UsdPhysics
shape_classes = (
("mesh", UsdGeom.Mesh),
("plane", UsdGeom.Plane),
("cube", UsdGeom.Cube),
("sphere", UsdGeom.Sphere),
("capsule", UsdGeom.Capsule),
("cylinder", UsdGeom.Cylinder),
("cone", UsdGeom.Cone),
)
shape_kind = next(
(name for name, shape_class in shape_classes if collision.IsA(shape_class)), "unknown"
)
result: dict[str, Any] = {
"schema_version": 1,
"path": str(collision.GetPath()),
"usd_authored_type_name": str(collision.GetTypeName()),
"usd_schema_type_name": str(collision.GetPrimTypeInfo().GetSchemaTypeName()),
"shape_kind": shape_kind,
"collision_enabled": bool(
UsdPhysics.CollisionAPI(collision).GetCollisionEnabledAttr().Get()
),
"physx_collision_api_present": bool(collision.HasAPI(PhysxSchema.PhysxCollisionAPI)),
}
if not result["usd_authored_type_name"] and not result["usd_schema_type_name"]:
raise RuntimeError(f"native collision {result['path']} has no observable schema type")
result.update(
{
name: None
for name in (
"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",
)
}
)
approximation: str | None = None
if collision.HasAPI(UsdPhysics.MeshCollisionAPI):
approximation_attr = UsdPhysics.MeshCollisionAPI(collision).GetApproximationAttr()
approximation = str(approximation_attr.Get()) if approximation_attr.Get() else None
result["mesh_approximation"] = approximation
physx_collision = PhysxSchema.PhysxCollisionAPI(collision)
for prefix, getter in (
("contact_offset", physx_collision.GetContactOffsetAttr),
("rest_offset", physx_collision.GetRestOffsetAttr),
):
if result["physx_collision_api_present"]:
attribute = getter()
effective = attribute.Get()
if isinstance(effective, bool) or not isinstance(effective, (int, float)):
raise RuntimeError(f"native {prefix} is malformed")
authored = attribute.Get() if attribute.HasAuthoredValue() else None
auto = bool(np.isneginf(effective))
result[prefix + "_authored"] = (
float(authored) if authored is not None and np.isfinite(authored) else None
)
result[prefix + "_schema_value"] = "simulation_determined" if auto else float(effective)
result[prefix + "_simulation_determined"] = auto
else:
result[prefix + "_authored"] = None
result[prefix + "_schema_value"] = None
result[prefix + "_simulation_determined"] = None
xformable = UsdGeom.Xformable(collision)
time = Usd.TimeCode.Default()
world_transform = xformable.ComputeLocalToWorldTransform(time)
parent_transform = xformable.ComputeParentToWorldTransform(time)
local_transform = np.matmul(
np.asarray(world_transform, dtype=np.float64),
np.linalg.inv(np.asarray(parent_transform, dtype=np.float64)),
)
result["world_transform"] = _finite_matrix(world_transform, "world transform")
result["local_transform"] = _finite_matrix(local_transform, "local transform")
local_box = bbox_cache.ComputeLocalBound(collision)
world_box = bbox_cache.ComputeWorldBound(collision)
local_empty, local_min, local_max = _range_bounds(
local_box.ComputeAlignedRange() if local_box else None, "local bounds"
)
world_empty, world_min, world_max = _range_bounds(
world_box.ComputeAlignedRange() if world_box else None, "world bounds"
)
local_extent = (
None
if local_min is None or local_max is None
else (np.asarray(local_max) - np.asarray(local_min)).tolist()
)
world_extent = (
None
if world_min is None or world_max is None
else (np.asarray(world_max) - np.asarray(world_min)).tolist()
)
extent_tolerance = 1e-12
world_axis_count = (
3 if world_extent is None else sum(value > extent_tolerance for value in world_extent)
)
zero_extent = bool(world_empty or (world_extent is not None and world_axis_count == 0))
zero_area = bool(world_empty or world_axis_count <= 1)
zero_volume_extent = bool(world_empty or world_axis_count <= 2)
planar_extent = bool(not world_empty and world_axis_count == 2)
result.update(
{
"local_bounds_empty": local_empty,
"local_bounds_min": local_min,
"local_bounds_max": local_max,
"local_extent": local_extent,
"world_bounds_empty": world_empty,
"world_bounds_min": world_min,
"world_bounds_max": world_max,
"world_extent": world_extent,
"extent_tolerance": extent_tolerance,
"zero_extent": zero_extent,
"zero_area": zero_area,
"zero_volume_extent": zero_volume_extent,
"planar_extent": planar_extent,
}
)
if shape_kind == "mesh":
result.update(_mesh_summary(UsdGeom.Mesh(collision)))
result["zero_extent"] = bool(result["zero_extent"] or result["mesh_points_zero_extent"])
result["zero_area"] = bool(result["zero_area"] or result["mesh_zero_area"])
result["zero_volume_extent"] = bool(
result["zero_volume_extent"] or result["mesh_zero_area"]
)
return result
def _normalized_variant_record(entity: Any, record: dict[str, Any]) -> dict[str, Any]:
"""Pad an optional-slot-omitting variant record to the full public layout.
``uniform_public_layout`` permits optional mesh-geom slots to be absent
from a variant (e.g. eraser heads), but PhysX views are homogeneous across
environments: every destination copy must carry the same native shape
count. The worker therefore normalizes each subset record to the compiled
public layout, marking padded rows through the worker-owned
``geom_placeholders`` field; ``_bake`` authors collision-disabled,
invisible placeholder prims for exactly those rows. Records already
covering the full layout pass through unchanged.
"""
if "geom_placeholders" in record:
raise ValueError("variant geom_placeholders are worker-owned")
public_geoms = [(geom.name, geom.body_name) for geom in entity.geoms]
actual_geoms = list(zip(record["geom_names"], record["geom_body_names"]))
if actual_geoms == public_geoms:
return record
slots: list[int] = []
cursor = 0
for pair in actual_geoms:
while cursor < len(public_geoms) and public_geoms[cursor] != pair:
cursor += 1
if cursor == len(public_geoms):
raise ValueError("variant geom names differ from compiled layout")
slots.append(cursor)
cursor += 1
present = set(slots)
contype = [0] * len(public_geoms)
conaffinity = [0] * len(public_geoms)
friction = [[0.0, 0.0, 0.0] for _ in public_geoms]
source_geom_types = _validated_variant_geom_types(record, len(actual_geoms))
geom_types = ["other"] * len(public_geoms)
_validated_variant_source_geometry(record, len(actual_geoms))
source_types = ["absent"] * len(public_geoms)
source_sizes = [[0.0, 0.0, 0.0] for _ in public_geoms]
source_poses = [[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0] for _ in public_geoms]
placeholders = [0] * len(public_geoms)
for source, slot in enumerate(slots):
contype[slot] = int(record["geom_contype"][source])
conaffinity[slot] = int(record["geom_conaffinity"][source])
friction[slot] = [float(value) for value in record["geom_friction"][source]]
geom_types[slot] = source_geom_types[source]
source_types[slot] = str(record["geom_source_types"][source])
source_sizes[slot] = [float(value) for value in record["geom_source_sizes"][source]]
source_poses[slot] = [float(value) for value in record["geom_source_poses"][source]]
for slot in range(len(public_geoms)):
if slot not in present:
placeholders[slot] = 1
normalized = dict(record)
normalized["geom_names"] = [name for name, _body in public_geoms]
normalized["geom_body_names"] = [body for _name, body in public_geoms]
normalized["geom_contype"] = contype
normalized["geom_conaffinity"] = conaffinity
normalized["geom_friction"] = friction
normalized["geom_types"] = geom_types
normalized["geom_source_types"] = source_types
normalized["geom_source_sizes"] = source_sizes
normalized["geom_source_poses"] = source_poses
normalized["geom_placeholders"] = placeholders
return normalized
def _consistent_native_mask(entry: dict[str, Any]) -> np.ndarray:
"""One native geometry layout shared by every variant, failing closed on drift."""
masks = {
tuple(bool(bit) for bit in _record_native_mask(record)) for record in entry["variants"]
}
if len(masks) != 1:
raise RuntimeError("entity variants disagree on the native geometry layout")
return np.asarray(masks.pop(), dtype=bool)
def _author_placeholder_geoms(
stage: Any,
root_path: str,
body_paths: dict[str, str],
record: dict[str, Any],
contact_offset: float | None = None,
rest_offset: float | None = None,
) -> None:
"""Author collision-disabled placeholder prims for variant-absent slots.
PhysX views are homogeneous across environments, so a variant omitting an
optional public slot still owns a native shape for it: a tiny invisible
sphere with collision disabled. The prim carries the slot's authored
identity (``unisim:geomIndex``/``unisim:geomName``), which
``_author_native_geometry`` and the readback match on, and the worker-only
``unisim:placeholder`` marker. Collision-disabled shapes contribute no
contacts and no mass; explicit body MassAPI attributes are unchanged.
"""
from pxr import PhysxSchema, Sdf, UsdGeom, UsdPhysics
placeholders = record.get("geom_placeholders")
if not placeholders or not any(placeholders):
return
for geom_index, is_placeholder in enumerate(placeholders):
if not is_placeholder:
continue
body_name = record["geom_body_names"][geom_index]
body_prim = stage.GetPrimAtPath(root_path + body_paths[body_name])
if not body_prim or not body_prim.IsValid():
raise RuntimeError(f"placeholder geom body prim is missing: {body_name}")
sphere = UsdGeom.Sphere.Define(
stage, f"{body_prim.GetPath()}/unisim_placeholder_geom_{geom_index}"
)
sphere.CreateRadiusAttr().Set(1e-4)
prim = sphere.GetPrim()
prim.CreateAttribute("unisim:geomName", Sdf.ValueTypeNames.String).Set(
record["geom_names"][geom_index]
)
prim.CreateAttribute("unisim:geomIndex", Sdf.ValueTypeNames.Int).Set(geom_index)
prim.CreateAttribute("unisim:placeholder", Sdf.ValueTypeNames.Bool).Set(True)
UsdPhysics.CollisionAPI.Apply(prim).CreateCollisionEnabledAttr().Set(False)
# Solver offsets are entity-level configuration; bake them onto
# placeholder prims too so cache-hit role inspection revalidates one
# uniform value set across every CollisionAPI prim.
if contact_offset is not None or rest_offset is not None:
physx_collision = PhysxSchema.PhysxCollisionAPI.Apply(prim)
if contact_offset is not None:
physx_collision.CreateContactOffsetAttr().Set(float(contact_offset))
if rest_offset is not None:
physx_collision.CreateRestOffsetAttr().Set(float(rest_offset))
UsdGeom.Imageable(prim).MakeInvisible()
def _author_native_geometry(
stage: Any,
root_path: str,
body_paths: dict[str, str],
entity: Any,
record: dict[str, Any],
*,
share_friction_materials: bool = False,
) -> None:
"""Author source-indexed collision identity and effective friction materials."""
from pxr import Sdf, UsdPhysics, UsdShade
# Visual-only geoms (contype=0 & conaffinity=0) carry no CollisionAPI in
# the converted USD (the importer honors the contact bits), so identity
# and friction are authored on the native subset only; ``geomIndex``
# keeps the public (full) record position, matching all-colliding scenes
# value-for-value. Placeholder prims arrive pre-slotted by
# ``_author_placeholder_geoms`` and are matched by that authored identity
# instead of by prim order, so placeholder placement is order-free.
native_mask = _record_native_mask(record)
authored = 0
shared_material_paths: dict[float, str] | None
if share_friction_materials:
shared_material_paths = {}
else:
shared_material_paths = None
for body_name in entity.body_names:
body_prim = stage.GetPrimAtPath(root_path + body_paths[body_name])
if not body_prim or not body_prim.IsValid():
raise RuntimeError(f"entity {entity.name} native body prim is missing: {body_name}")
expected = [
index
for index, owner in enumerate(record["geom_body_names"])
if owner == body_name and native_mask[index]
]
collisions = _body_collision_prims(body_prim)
if len(collisions) != len(expected):
raise RuntimeError(
f"entity {entity.name} body {body_name} has "
f"{len(collisions)} native collision geoms, expected {len(expected)}"
)
preslotted: dict[int, Any] = {}
unmarked: list[Any] = []
for collision in collisions:
if collision.HasAttribute("unisim:geomIndex"):
observed = collision.GetAttribute("unisim:geomIndex").Get()
if isinstance(observed, bool) or not isinstance(observed, int):
raise RuntimeError(f"entity {entity.name} placeholder geom index is invalid")
if observed in preslotted:
raise RuntimeError(f"entity {entity.name} placeholder geom index is duplicated")
preslotted[observed] = collision
else:
unmarked.append(collision)
remaining = [index for index in expected if index not in preslotted]
if len(unmarked) != len(remaining) or any(index not in expected for index in preslotted):
raise RuntimeError(
f"entity {entity.name} body {body_name} native geometry differs from source"
)
for geom_index, collision in zip(remaining, unmarked):
collision.CreateAttribute("unisim:geomName", Sdf.ValueTypeNames.String).Set(
record["geom_names"][geom_index]
)
collision.CreateAttribute("unisim:geomIndex", Sdf.ValueTypeNames.Int).Set(geom_index)
preslotted[geom_index] = collision
for geom_index in expected:
collision = preslotted[geom_index]
observed_name = collision.GetAttribute("unisim:geomName").Get()
if observed_name != record["geom_names"][geom_index]:
raise RuntimeError(f"entity {entity.name} native geometry identity differs")
sliding_friction = float(record["geom_friction"][geom_index][0])
material_path = None
if shared_material_paths is not None:
material_path = shared_material_paths.get(sliding_friction)
if material_path is None:
name = (
f"unisim_friction_{len(shared_material_paths)}"
if shared_material_paths is not None
else f"unisim_geom_{geom_index}"
)
material_path = f"{root_path}/Looks/{name}"
material = UsdShade.Material.Define(stage, material_path)
physics_material = UsdPhysics.MaterialAPI.Apply(material.GetPrim())
physics_material.CreateStaticFrictionAttr().Set(sliding_friction)
physics_material.CreateDynamicFrictionAttr().Set(sliding_friction)
if shared_material_paths is not None:
shared_material_paths[sliding_friction] = material_path
collision.CreateRelationship("material:binding:physics").SetTargets(
[Sdf.Path(material_path)]
)
authored += 1
if authored != int(native_mask.sum()):
raise RuntimeError(f"entity {entity.name} native geometry record is incomplete")
def _native_geometry_record(
root: Any, entity: Any, record: dict[str, Any], body_paths: dict[str, str]
) -> dict[str, Any]:
"""Read geometry identity/contact/friction from one actual native subtree."""
from pxr import Usd, UsdGeom, UsdPhysics
root_path = str(root.GetPath())
names: list[str] = []
body_names: list[str] = []
masks: list[list[int]] = []
friction: list[list[float]] = []
source_types = [str(value) for value in record["geom_source_types"]]
source_sizes = [[float(value) for value in row] for row in record["geom_source_sizes"]]
source_poses = [[float(value) for value in row] for row in record["geom_source_poses"]]
native_audits: list[dict[str, Any] | None] = [None for _ in record["geom_names"]]
bbox_cache = UsdGeom.BBoxCache(
Usd.TimeCode.Default(), [UsdGeom.Tokens.default_], useExtentsHint=False
)
# Visual-only geoms have no native collision prim (see
# _author_native_geometry); they keep their public row with a zero mask
# and the record's friction, and only the native subset is read back.
# Native prims are matched through their authored ``unisim:geomIndex``
# identity, so placeholder placement never depends on prim order.
native_mask = _record_native_mask(record)
for body_name in entity.body_names:
body_prim = root.GetStage().GetPrimAtPath(root_path + body_paths[body_name])
if not body_prim or not body_prim.IsValid():
raise RuntimeError(f"entity {entity.name} native body prim is missing: {body_name}")
expected = [
index
for index, owner in enumerate(record["geom_body_names"])
if owner == body_name and native_mask[index]
]
collisions = _body_collision_prims(body_prim)
if len(collisions) != len(expected):
raise RuntimeError(
f"entity {entity.name} body {body_name} native geometry differs from source"
)
by_index: dict[int, Any] = {}
for collision in collisions:
observed_index = collision.GetAttribute("unisim:geomIndex").Get()
if (
isinstance(observed_index, bool)
or not isinstance(observed_index, int)
or observed_index in by_index
):
raise RuntimeError(f"entity {entity.name} native geometry identity differs")
by_index[observed_index] = collision
for geom_index in expected:
collision = by_index.get(geom_index)
if collision is None:
raise RuntimeError(f"entity {entity.name} native geometry identity differs")
observed_name = collision.GetAttribute("unisim:geomName").Get()
if observed_name != record["geom_names"][geom_index]:
raise RuntimeError(f"entity {entity.name} native geometry identity differs")
enabled = UsdPhysics.CollisionAPI(collision).GetCollisionEnabledAttr().Get()
if not isinstance(enabled, bool):
raise RuntimeError(f"entity {entity.name} native collision state is missing")
bindings = collision.GetRelationship("material:binding:physics").GetTargets()
if len(bindings) != 1:
raise RuntimeError(f"entity {entity.name} native geometry material is missing")
material = root.GetStage().GetPrimAtPath(bindings[0])
if (
not material
or not material.IsValid()
or not material.HasAPI(UsdPhysics.MaterialAPI)
):
raise RuntimeError(f"entity {entity.name} native geometry material is invalid")
physics_material = UsdPhysics.MaterialAPI(material)
static_friction = physics_material.GetStaticFrictionAttr().Get()
dynamic_friction = physics_material.GetDynamicFrictionAttr().Get()
values = [static_friction, dynamic_friction, 0.0]
if (
not all(isinstance(value, (int, float)) for value in values[:2])
or not np.isfinite(values).all()
or any(value < 0.0 for value in values)
):
raise RuntimeError(f"entity {entity.name} native geometry friction is invalid")
while len(names) < geom_index:
# Visual-only rows keep the public layout with a zero mask and
# the record's friction (no native shape exists to read back).
skipped = len(names)
names.append(str(record["geom_names"][skipped]))
body_names.append(str(record["geom_body_names"][skipped]))
masks.append([0, 0])
friction.append([float(v) for v in record["geom_friction"][skipped]])
native_audits[skipped] = None
names.append(str(observed_name))
body_names.append(body_name)
masks.append([int(enabled), int(enabled)])
friction.append([float(value) for value in values])
native_audits[geom_index] = _native_collision_geometry_audit(collision, bbox_cache)
while len(names) < len(record["geom_names"]):
skipped = len(names)
names.append(str(record["geom_names"][skipped]))
body_names.append(str(record["geom_body_names"][skipped]))
masks.append([0, 0])
friction.append([float(v) for v in record["geom_friction"][skipped]])
if names != list(record["geom_names"]) or body_names != list(record["geom_body_names"]):
raise RuntimeError(f"entity {entity.name} native geometry layout differs from source")
return {
"geom_names": names,
"geom_body_names": body_names,
"geom_contact_masks": masks,
"geom_friction": friction,
"geometry_audit_schema_version": 1,
"geom_source_types": source_types,
"geom_source_sizes": source_sizes,
"geom_source_poses": source_poses,
"geom_native_audits": native_audits,
}
def _inspect_role(
usd_path: str,
entity: Any,
entry: dict[str, Any],
variant: int,
require_bodies: bool = False,
contact_offset: float | None = None,
rest_offset: float | None = None,
) -> tuple[str, dict[str, str]]:
"""Read and validate an immutable role artifact without authoring edits."""
from pxr import PhysxSchema, Usd, UsdPhysics
stage = Usd.Stage.Open(usd_path)
root = stage.GetDefaultPrim()
if not root or not root.IsValid():
raise RuntimeError("role entity USD has no default prim")
observed_variant = root.GetAttribute("unisim:variantIndex").Get()
if observed_variant != variant:
raise RuntimeError("role entity USD variant identity differs from request")
body_paths: dict[str, str] = {}
articulation_roots: list[str] = []
active_joints = 0
rigid_body_count = 0
root_path = str(root.GetPath())
expected_collision = bool(entry["collision_enabled"])
expected_disable_gravity = bool(entry["gravity_disabled"])
for prim in Usd.PrimRange(root):
if prim.IsA(UsdPhysics.Joint) and prim.IsActive():
active_joints += 1
if prim.HasAPI(UsdPhysics.ArticulationRootAPI):
articulation_roots.append(str(prim.GetPath()))
if prim.HasAPI(UsdPhysics.RigidBodyAPI):
rigid_body_count += 1
body_name = prim.GetName()
prim_path = str(prim.GetPath())
relative_body_path = prim_path[len(root_path) :]
if body_name in body_paths and body_paths[body_name] != relative_body_path:
raise RuntimeError(f"entity {entity.name} has duplicate rigid body {body_name!r}")
body_paths[body_name] = relative_body_path
if entity.kind == "rigid":
kinematic = UsdPhysics.RigidBodyAPI(prim).GetKinematicEnabledAttr().Get()
if kinematic is not (entity.root_mode != "floating"):
raise RuntimeError(f"rigid entity {entity.name} has the wrong mobility role")
disable_gravity = PhysxSchema.PhysxRigidBodyAPI(prim).GetDisableGravityAttr().Get()
if disable_gravity is not expected_disable_gravity:
raise RuntimeError(f"entity {entity.name} has the wrong gravity role")
if require_bodies:
sleep_threshold = PhysxSchema.PhysxRigidBodyAPI(prim).GetSleepThresholdAttr().Get()
report_threshold = PhysxSchema.PhysxContactReportAPI(prim).GetThresholdAttr().Get()
if sleep_threshold != 0.0 or report_threshold != 0.0:
raise RuntimeError(f"entity {entity.name} has the wrong contact-reporting role")
if prim.HasAPI(UsdPhysics.CollisionAPI):
collision = UsdPhysics.CollisionAPI(prim).GetCollisionEnabledAttr().Get()
placeholder = prim.HasAttribute("unisim:placeholder") and (
prim.GetAttribute("unisim:placeholder").Get() is True
)
if placeholder:
# Placeholder prims stand in for variant-absent optional
# slots; they must never collide.
if collision is not False:
raise RuntimeError(f"entity {entity.name} has an enabled placeholder geom")
elif collision is not expected_collision:
raise RuntimeError(f"entity {entity.name} has the wrong collision role")
# The cache identity tracks the requested offsets; revalidate the
# authored values like every other baked role semantic. PhysX
# stores these attributes in single precision, so compare against
# the float32 rounding of the request.
if contact_offset is not None:
authored = PhysxSchema.PhysxCollisionAPI(prim).GetContactOffsetAttr().Get()
if authored != float(np.float32(contact_offset)):
raise RuntimeError(f"entity {entity.name} has the wrong baked contact offset")
if rest_offset is not None:
authored = PhysxSchema.PhysxCollisionAPI(prim).GetRestOffsetAttr().Get()
if authored != float(np.float32(rest_offset)):
raise RuntimeError(f"entity {entity.name} has the wrong baked rest offset")
for axis in ("angular", "linear"):
if prim.HasAPI(UsdPhysics.DriveAPI, axis):
drive = UsdPhysics.DriveAPI(prim, axis)
if drive.GetStiffnessAttr().Get() != 0.0 or (drive.GetDampingAttr().Get() != 0.0):
raise RuntimeError(f"entity {entity.name} retains a converter drive")
if entity.kind == "rigid":
if rigid_body_count != 1:
raise RuntimeError(f"rigid entity {entity.name} has {rigid_body_count} native bodies")
if active_joints:
raise RuntimeError(f"rigid entity {entity.name} retains native joints")
if articulation_roots:
raise RuntimeError(f"rigid entity {entity.name} retains an articulation root")
relative = ""
else:
if len(articulation_roots) != 1:
raise RuntimeError(
f"entity {entity.name} has ambiguous articulation roots: {articulation_roots}"
)
articulation_root = articulation_roots[0]
if entity.root_mode == "fixed":
if articulation_root != root_path:
raise RuntimeError(f"fixed entity {entity.name} has no root-prim articulation")
relative = ""
else:
if stage.GetPrimAtPath(articulation_root).GetName() != entity.root_body:
raise RuntimeError(
f"entity {entity.name} articulation root differs from declaration"
)
relative = articulation_root[len(root_path) :]
flag = PhysxSchema.PhysxArticulationAPI(
stage.GetPrimAtPath(articulation_root)
).GetEnabledSelfCollisionsAttr()
if bool(flag.Get()) != bool(entry["self_collision"]):
raise RuntimeError(f"entity {entity.name} has the wrong self-collision role")
missing = [name for name in entity.body_names if name not in body_paths]
if require_bodies and missing:
raise RuntimeError(f"entity {entity.name} role USD is missing bodies: {missing}")
_native_geometry_record(root, entity, entry["variants"][variant], body_paths)
return relative, body_paths
[docs]
class SceneWorkerContext:
"""One independent native asset view per declared entity."""
[docs]
def __init__(self, protocol: Any, renderer: Any) -> None:
self.protocol = protocol
self.renderer = renderer
self.progress_out: Any = None
self.slots: dict[str, np.ndarray] = {}
self._shm_handles: list[Any] = []
self.assets: list[Any] = []
self.maps: list[dict[str, Any]] = []
self.contact_sensors: list[Any] = []
self.contact_sensor_maps: list[dict[str, Any]] = []
self.contact_force_sensors: list[dict[str, str]] = []
self.net_contact_entities: list[str] = []
self.net_contact_views: list[Any] = []
self.net_contact_maps: list[dict[str, Any]] = []
self._contact_reporting = False
self._cuda_ipc: WorkerCudaIpcArena | None = None
self._cuda_maps: list[dict[str, Any]] = []
self._cuda_sensor_specs: dict[str, dict[str, Any]] = {}
self._cuda_origins: Any = None
self._cuda_reset_sequence = 0
self._tensor_cuda_ipc = False
self._share_friction_materials = False
self._env_spacing = 2.0
self.faulted = False
self.model_file_projection: Any = None
self._model_file_metadata: dict[str, Any] | None = None
self._init_telemetry: dict[str, Any] | None = None
# Per-entity drive/natural damping bookkeeping in public joint order.
# PhysX exposes one DOF damping quantity per joint; the implicit
# position drive's kd and MuJoCo-style natural joint damping compose
# additively (the drive velocity target is always zero), so the worker
# tracks both channels and writes their sum.
self._drive_damping: list[np.ndarray | None] = []
self._natural_damping: list[np.ndarray | None] = []
self.physx_solver = PhysxSolverConfig()
self._raw_usd_cache: RawUSDCache | None = None
self._raw_usd_cache_persistent = False
self._raw_usd_cache_reports: list[dict[str, Any]] = []
self._reported_raw_usd_identities: set[str] = set()
self._reported_raw_usd_source_digests: set[str] = set()
self._role_usd_cache: RoleUSDCache | None = None
self._role_usd_cache_reports: list[dict[str, Any]] = []
self._temporary = tempfile.TemporaryDirectory(prefix="unisim-isaacsim-scene-")
@staticmethod
def _native_visual_sphere_radii(prim: Any) -> list[float]:
"""Read sphere dimensions from the body's imported visual subtree."""
from pxr import Usd, UsdGeom
visuals = prim.GetChild("visuals")
if not visuals or not visuals.IsValid():
return []
result: list[float] = []
for child in Usd.PrimRange(visuals):
if not child.IsA(UsdGeom.Sphere):
continue
radius = UsdGeom.Sphere(child).GetRadiusAttr().Get()
if radius is None:
raise RuntimeError(f"native sphere {child.GetPath()} has no radius")
value = float(radius)
if not np.isfinite(value) or value <= 0.0:
raise RuntimeError(f"native sphere {child.GetPath()} has an invalid radius")
result.append(value)
return result
@staticmethod
def _native_source_sphere_radii(
prim: Any, record: dict[str, Any], body_name: str
) -> list[float]:
"""Read source-owned sphere radii in public geometry order.
The USD importer can represent capsule colliders with additional sphere
approximation prims. Those shapes carry the capsule's source geom index,
so use compiled geometry-type identity rather than counting every USD
sphere below a rigid body.
"""
from pxr import Usd, UsdGeom
geom_types = _validated_variant_geom_types(record, len(record["geom_body_names"]))
expected = [
index
for index, (owner, kind) in enumerate(zip(record["geom_body_names"], geom_types))
if owner == body_name and kind == "sphere"
]
collision_mask = _record_collision_mask(record)
indexed: dict[int, float] = {}
visual: list[float] = []
visuals = prim.GetChild("visuals")
visuals_path = str(visuals.GetPath()) if visuals and visuals.IsValid() else None
for child in Usd.PrimRange(prim):
if not child.IsA(UsdGeom.Sphere):
continue
placeholder = child.GetAttribute("unisim:placeholder")
if placeholder.IsDefined() and placeholder.Get() is True:
continue
radius = UsdGeom.Sphere(child).GetRadiusAttr().Get()
if radius is None:
raise RuntimeError(f"native sphere {child.GetPath()} has no radius")
value = float(radius)
if not np.isfinite(value) or value <= 0.0:
raise RuntimeError(f"native sphere {child.GetPath()} has an invalid radius")
observed_index = child.GetAttribute("unisim:geomIndex").Get()
if (
isinstance(observed_index, int)
and not isinstance(observed_index, bool)
and observed_index in expected
):
indexed[observed_index] = value
elif visuals_path is not None and str(child.GetPath()).startswith(visuals_path + "/"):
visual.append(value)
result: list[float] = []
visual_cursor = 0
for index in expected:
if bool(collision_mask[index]):
result.append(indexed[index])
elif visual_cursor < len(visual):
result.append(visual[visual_cursor])
visual_cursor += 1
else:
result.append(float("nan"))
return result
def _tensor(self, values: np.ndarray) -> Any:
return self.torch.as_tensor(
np.ascontiguousarray(values), dtype=self.torch.float32, device=self.device
)
def _cpu_tensor(self, values: np.ndarray) -> Any:
return self.torch.as_tensor(
np.ascontiguousarray(values), dtype=self.torch.float32, device="cpu"
)
[docs]
def init_sim(self, payload: dict[str, Any]) -> dict[str, Any]:
if "scene_layout" not in payload:
metadata = self.renderer.init_sim(payload)
self.adopt_initialized_context(self.renderer, metadata, payload)
return cast(dict[str, Any], metadata)
self.layout = validate_scene_payload(self.protocol, payload)
self._tensor_cuda_ipc = bool(payload.get("tensor_cuda_ipc"))
share_friction_materials = payload.get("share_friction_materials", False)
if not isinstance(share_friction_materials, bool):
raise ValueError("share_friction_materials must be a boolean")
self._share_friction_materials = share_friction_materials
telemetry = _InitTelemetry()
self.contact_force_sensors = _validate_contact_force_sensors(payload, self.layout)
self.net_contact_entities = _validate_body_net_contact_entities(payload, self.layout)
self._contact_reporting = bool(self.contact_force_sensors) or bool(
self.net_contact_entities
)
raw_usd_cache_dir = payload.get("raw_usd_cache_dir")
if raw_usd_cache_dir is not None and (
not isinstance(raw_usd_cache_dir, str) or not raw_usd_cache_dir
):
raise ValueError("raw USD cache directory must be a non-empty string or null")
role_usd_cache_dir = payload.get("role_usd_cache_dir")
if role_usd_cache_dir is not None and (
not isinstance(role_usd_cache_dir, str) or not role_usd_cache_dir
):
raise ValueError("role USD cache directory must be a non-empty string or null")
self._raw_usd_cache_persistent = raw_usd_cache_dir is not None
self._raw_usd_cache = RawUSDCache(
Path(raw_usd_cache_dir)
if raw_usd_cache_dir is not None
else Path(self._temporary.name) / "raw-usd"
)
if role_usd_cache_dir is not None:
self._role_usd_cache = RoleUSDCache(Path(role_usd_cache_dir))
self.entity_components = {
entity.name: _entity_prim_component(entity.name) for entity in self.layout.entities
}
self.num_envs = payload["num_envs"]
self.entries = payload["scene_entities"]
self.sim_dt = float(payload["sim_dt"])
env_spacing = payload.get("env_spacing", 2.0)
if (
isinstance(env_spacing, bool)
or not isinstance(env_spacing, (int, float))
or 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)
device_id = int(payload.get("device_id", 0))
if device_id < 0:
raise NotImplementedError("IsaacSim mapped scene requires CUDA")
self.device = f"cuda:{device_id}"
render_mode = payload.get("render_mode", "none")
if render_mode not in ("none", "record", "interactive"):
raise ValueError("invalid render_mode")
os.environ["HEADLESS"] = "0" if render_mode == "interactive" else "1"
os.environ["ENABLE_CAMERAS"] = "1" if render_mode == "record" else "0"
os.environ["LIVESTREAM"] = "0"
os.environ["XR"] = "0"
progress = self.protocol.InitProgress(self, enabled=bool(payload.get("init_progress")))
progress.report("isaacsim worker: starting kit", 0, 1, force=True)
os.environ.setdefault("OMNI_KIT_ACCEPT_EULA", "1")
from isaaclab.app import AppLauncher
app = AppLauncher(
{
"headless": render_mode != "interactive",
"enable_cameras": render_mode == "record",
"device": self.device,
"multi_gpu": False,
}
).app
self.renderer.simulation_app = app
import isaaclab.sim as sim_utils
import isaacsim.core.utils.prims as prim_utils
import torch
from isaaclab.actuators import ImplicitActuatorCfg
from isaaclab.assets import Articulation, ArticulationCfg, RigidObject, RigidObjectCfg
from isaaclab.sensors import ContactSensor, ContactSensorCfg
from isaaclab.sim.converters import MjcfConverter, MjcfConverterCfg
from isaacsim.core.cloner import Cloner, GridCloner
from isaacsim.core.utils.extensions import enable_extension
from pxr import Sdf
self.torch = torch
enable_extension("isaacsim.asset.importer.mjcf")
self.gravity = np.asarray(payload["gravity"], dtype=np.float64)
if self.gravity.shape != (3,) or not np.isfinite(self.gravity).all():
raise ValueError("invalid gravity")
self.physx_solver = PhysxSolverConfig.from_payload(payload.get("physx_solver"))
self.sim = sim_utils.SimulationContext(
sim_utils.SimulationCfg(
dt=self.sim_dt,
device=self.device,
gravity=tuple(self.gravity.tolist()),
physx=build_isaaclab_physx_cfg(sim_utils, self.physx_solver),
)
)
telemetry.mark("kit_startup_sim_context")
progress.report("isaacsim worker: starting kit", 1, 1, force=True)
if self._contact_reporting:
# IsaacLab disables PhysX contact processing by default and only
# ContactSensor construction re-enables it; the body-net path uses
# raw PhysX contact views, so enable reporting explicitly.
import carb
carb.settings.get_settings().set_bool("/physics/disableContactProcessing", False)
cloner = GridCloner(spacing=self._env_spacing)
prototype_cloner = Cloner()
cloner.define_base_env("/World/envs")
self.env_paths = cloner.generate_paths("/World/envs/env", self.num_envs)
self.entity_paths = {
name: [path + "/" + component for path in self.env_paths]
for name, component in self.entity_components.items()
}
prim_utils.create_prim(self.env_paths[0], "Xform")
self.origins = np.asarray(
cloner.clone(
source_prim_path=self.env_paths[0],
prim_paths=self.env_paths,
replicate_physics=False,
copy_from_source=True,
),
dtype=np.float32,
)
telemetry.mark("env_grid_clone")
telemetry.count_prims(self.sim.stage, "after_env_grid")
self.usd_paths = []
self.entity_body_paths: list[list[dict[str, str]]] = []
content_identity = SceneContentIdentity.from_dict(payload["scene_content_identity"])
raw_usd_runtime_versions = _raw_usd_runtime_versions()
layout_entities = {entity.name: entity for entity in self.layout.entities}
layout_entries = {entry["name"]: entry for entry in self.entries}
# One counted unit per variant materialization, per prototype spawn
# and per destination copy keeps a single monotonic bar across the
# per-entity loops below.
build_total = 2 * sum(len(entry["sources"]) for entry in self.entries)
build_total += self.num_envs * len(self.entries)
build_done = 0
for entity, entry in zip(self.layout.entities, self.entries):
entity_started = time.perf_counter()
component = self.entity_components[entity.name]
paths, root_paths = [], []
body_paths_by_variant: list[dict[str, str]] = []
for index, source in enumerate(entry["sources"]):
assert self._raw_usd_cache is not None
raw_cache = self._raw_usd_cache
raw_entity = (
entity if entry["mirror_of"] is None else layout_entities[entry["mirror_of"]]
)
# Raw conversion is role-neutral: a mirror shares its source
# entity's raw USD, including the source's self-collision flag.
raw_self_collision = bool(layout_entries[raw_entity.name]["self_collision"])
raw_request = _raw_usd_request(
content_identity,
source,
raw_entity,
index,
raw_usd_runtime_versions,
self_collision=raw_self_collision,
)
def convert_raw_usd(artifact_dir: Path, usd_file_name: str) -> Path:
raw_converter = MjcfConverter(
MjcfConverterCfg(
asset_path=source,
fix_base=(
raw_entity.kind == "articulation"
and raw_entity.root_mode == "fixed"
),
import_sites=False,
import_inertia_tensor=True,
make_instanceable=False,
self_collision=raw_self_collision,
force_usd_conversion=True,
usd_dir=str(artifact_dir),
usd_file_name=usd_file_name,
)
)
return Path(raw_converter.usd_path)
cached_raw = raw_cache.materialize(raw_request, convert_raw_usd)
telemetry.mark(f"entity.{entity.name}.raw_usd_cache_convert")
if cached_raw.record.identity not in self._reported_raw_usd_identities:
self._reported_raw_usd_identities.add(cached_raw.record.identity)
self._reported_raw_usd_source_digests.add(cached_raw.record.source_digest)
self._raw_usd_cache_reports.append(
{
"identity": cached_raw.record.identity,
"source_digest": cached_raw.record.source_digest,
"entity": raw_entity.name,
"variant": index,
"hit": cached_raw.hit,
"materialize_ms": cached_raw.materialize_ms,
"artifact_files": len(cached_raw.record.files),
"artifact_bytes": cached_raw.record.size_bytes,
"runtime_versions": dict(cached_raw.record.runtime_versions),
}
)
role_request = _role_usd_request(
cached_raw.record,
entity,
entry,
index,
require_bodies=self._contact_reporting,
contact_offset=self.physx_solver.contact_offset,
rest_offset=self.physx_solver.rest_offset,
share_friction_materials=self._share_friction_materials,
)
if self._role_usd_cache is None:
role_destination = Path(self._temporary.name) / "roles" / component / str(index)
usd_path = raw_cache.copy_artifact(cached_raw.record, role_destination)
body_paths: dict[str, str] = {}
relative = _bake(
str(usd_path),
entity,
entry,
index,
body_paths,
require_bodies=self._contact_reporting,
contact_offset=self.physx_solver.contact_offset,
rest_offset=self.physx_solver.rest_offset,
share_friction_materials=self._share_friction_materials,
)
self._role_usd_cache_reports.append(
{
"identity": role_request.identity,
"source_digest": cached_raw.record.source_digest,
"entity": entity.name,
"variant": index,
"hit": False,
"materialize_ms": 0.0,
"artifact_files": 0,
"artifact_bytes": 0,
"runtime_versions": dict(cached_raw.record.runtime_versions),
}
)
else:
if raw_cache.load(cached_raw.record.identity) != cached_raw.record:
raise RuntimeError("raw USD cache entry changed before role baking")
baked_body_paths: dict[str, str] | None = None
baked_root_path: str | None = None
def bake_role_artifact(artifact_dir: Path, _usd_file_name: str) -> Path:
nonlocal baked_body_paths, baked_root_path
raw_destination = artifact_dir.parent / "raw"
raw_cache.copy_artifact(cached_raw.record, raw_destination)
artifact_dir.rmdir()
os.replace(raw_destination, artifact_dir)
copied_usd = artifact_dir / Path(
*PurePosixPath(cached_raw.record.usd_relative_path).parts
)
body_paths: dict[str, str] = {}
baked_root_path = _bake(
str(copied_usd),
entity,
entry,
index,
body_paths,
require_bodies=self._contact_reporting,
contact_offset=self.physx_solver.contact_offset,
rest_offset=self.physx_solver.rest_offset,
share_friction_materials=self._share_friction_materials,
)
baked_body_paths = body_paths
return copied_usd
cached_role = self._role_usd_cache.materialize(role_request, bake_role_artifact)
if cached_role.hit:
relative, body_paths = _inspect_role(
str(cached_role.record.usd_path),
entity,
entry,
index,
require_bodies=self._contact_reporting,
contact_offset=self.physx_solver.contact_offset,
rest_offset=self.physx_solver.rest_offset,
)
else:
assert baked_body_paths is not None and baked_root_path is not None
body_paths = baked_body_paths
relative = baked_root_path
usd_path = cached_role.record.usd_path
self._role_usd_cache_reports.append(
{
"identity": cached_role.record.identity,
"source_digest": cached_raw.record.source_digest,
"entity": entity.name,
"variant": index,
"hit": cached_role.hit,
"materialize_ms": cached_role.materialize_ms,
"artifact_files": len(cached_role.record.files),
"artifact_bytes": cached_role.record.size_bytes,
"runtime_versions": dict(cached_raw.record.runtime_versions),
}
)
paths.append(str(usd_path))
root_paths.append(relative)
body_paths_by_variant.append(body_paths)
telemetry.mark(f"entity.{entity.name}.role_bake_or_inspect")
build_done += 1
progress.report("isaacsim worker: building entities", build_done, build_total)
if len(set(root_paths)) != 1:
raise RuntimeError("variant articulation root paths differ")
if self._contact_reporting:
canonical_body_paths = body_paths_by_variant[0]
for body_name in entity.body_names:
for variant_body_paths in body_paths_by_variant[1:]:
if variant_body_paths.get(body_name) != canonical_body_paths[body_name]:
raise RuntimeError(
f"entity {entity.name} variant rigid-body prim paths differ"
)
self.entity_body_paths.append(body_paths_by_variant)
self.usd_paths.append(paths)
prototype_paths = _prototype_spawn_paths(component, len(paths))
assignment = _validated_assignment(entry, self.num_envs)
destination_groups = _assignment_groups(
assignment, len(paths), self.entity_paths[entity.name]
)
prim_path = "/World/envs/env_.*/" + component
prim_utils.create_prim(f"/World/unisim_prototypes/{component}", "Scope")
# Prototype spawns, destination copies and the prototype scope
# removal all run under one change-listener toggle: physics never
# parses the transient prototype subtrees at all, only the final
# per-environment destinations it observes when the listener is
# re-enabled. Per-spawn physics change processing otherwise
# dominates INIT at catalog scale. Stage-object edits (create_prim)
# must stay outside an Sdf.ChangeBlock — deferred layer notices
# break DefinePrim recomposition — so only the pure-spec copies
# keep the block (Cloner.clone additionally repeats the toggle and
# one Sdf.ChangeBlock per call).
prototype_cloner.disable_change_listener()
try:
for prototype_path, prototype_usd_path in zip(prototype_paths, paths):
# Spawn the reference arc directly instead of through
# IsaacLab's UsdFileCfg spawner: spawn_from_usd first
# runs an asyncio omni.client stat of the USD path whose
# event loop yields a fixed ~0.1s sleep per call, which
# dominated INIT at catalog scale (127s of 157s at 1200
# variants). Every path here is a validated cache
# artifact (manifest + SHA-256) or a fresh copy of one,
# and all spawner overrides the worker used are None or
# unconsumed by the from-files spawner, so the spawner
# reduced to exactly this create_prim call.
if not os.path.isfile(prototype_usd_path):
raise RuntimeError(f"IsaacSim role USD is missing: {prototype_usd_path}")
prim_utils.create_prim(
prototype_path,
usd_path=prototype_usd_path,
translation=tuple(entry["initial_pose"][:3]),
orientation=tuple(entry["initial_pose"][3:]),
)
telemetry.mark(f"entity.{entity.name}.prototype_authoring")
build_done += 1
progress.report("isaacsim worker: building entities", build_done, build_total)
if any(destination_groups):
with Sdf.ChangeBlock():
for prototype_path, destinations in zip(
prototype_paths, destination_groups
):
if destinations:
_copy_prims_from_source(
self.sim.stage, prototype_path, destinations
)
build_done += len(destinations)
progress.report(
"isaacsim worker: building entities",
build_done,
build_total,
force=build_done == build_total,
)
telemetry.mark(f"entity.{entity.name}.batched_copy")
for prototype_path in prototype_paths:
prototype = prim_utils.get_prim_at_path(prototype_path)
if not prototype or not prototype.IsValid():
raise RuntimeError(f"IsaacSim prototype is missing: {prototype_path}")
# Destinations are full copies (copy_from_source), so the
# prototypes carry no live state once copied. Remove the whole
# prototype scope instead of deactivating it: deactivated
# prototypes keep every prim resident, and the growing stage
# makes each subsequent copy slower (superlinear INIT at 10k+
# envs).
self.sim.stage.RemovePrim(f"/World/unisim_prototypes/{component}")
telemetry.mark(f"entity.{entity.name}.prototype_scope_removal")
finally:
prototype_cloner.enable_change_listener()
if entity.kind == "articulation":
names = [joint.name for joint in entity.joints]
gains = self.renderer._actuator_dicts(entry["variants"][0], names)
actuators = {}
if names:
actuators["declared"] = ImplicitActuatorCfg(
joint_names_expr=names,
stiffness=gains["stiffness"],
damping=gains["damping"],
effort_limit_sim=gains["effort"],
velocity_limit_sim=gains["velocity_limit"],
armature=gains["armature"],
friction=gains["friction"],
)
asset = Articulation(
ArticulationCfg(
prim_path=prim_path,
articulation_root_prim_path=root_paths[0],
init_state=ArticulationCfg.InitialStateCfg(
pos=tuple(entry["initial_pose"][:3]),
rot=tuple(entry["initial_pose"][3:]),
),
spawn=None,
actuators=actuators,
)
)
else:
asset = RigidObject(
RigidObjectCfg(
prim_path=prim_path,
spawn=None,
init_state=RigidObjectCfg.InitialStateCfg(
pos=tuple(entry["initial_pose"][:3]),
rot=tuple(entry["initial_pose"][3:]),
),
)
)
self.assets.append(asset)
telemetry.span(f"entity.{entity.name}.total", entity_started)
telemetry.count_prims(self.sim.stage, f"after_entity_{entity.name}")
progress.report("isaacsim worker: finalizing", 0, 1, force=True)
entity_indexes = {entity.name: index for index, entity in enumerate(self.layout.entities)}
for record in self.contact_force_sensors:
source_index = entity_indexes[record["source_entity"]]
target_index = entity_indexes[record["target_entity"]]
source_component = self.entity_components[record["source_entity"]]
target_component = self.entity_components[record["target_entity"]]
source_path = (
"/World/envs/env_.*/"
+ source_component
+ self.entity_body_paths[source_index][0][record["source_body"]]
)
target_path = (
"/World/envs/env_.*/"
+ target_component
+ self.entity_body_paths[target_index][0][record["target_body"]]
)
self.contact_sensors.append(
ContactSensor(
ContactSensorCfg(
prim_path=source_path,
filter_prim_paths_expr=[target_path],
history_length=0,
)
)
)
self._bind_fixed_anchors()
if self.num_envs > 1:
cloner.filter_collisions(
self.sim.cfg.physics_prim_path, "/World/collisions", self.env_paths
)
telemetry.mark("filter_collisions")
self._setup_renderer(sim_utils, payload)
# Collision offsets are baked into the role artifacts (identical on
# every environment copy); the max depenetration velocity is still
# authored on every rigid body before the first physics step.
if self.physx_solver.max_depenetration_velocity is not None:
apply_max_depenetration_velocity(
self.sim.stage, self.physx_solver.max_depenetration_velocity
)
telemetry.mark("collision_offsets")
self.sim.reset()
for entity, asset in zip(self.layout.entities, self.assets):
asset.update(self.sim_dt)
for sensor, record in zip(self.contact_sensors, self.contact_force_sensors):
sensor.update(self.sim_dt)
native_paths = list(sensor.body_physx_view.prim_paths)
native_envs = _native_environment_order(
native_paths, self.entity_paths[record["source_entity"]]
)
self.contact_sensor_maps.append(
{"envs": np.argsort(native_envs), "public_for_native": np.asarray(native_envs)}
)
layout_entities = {entity.name: entity for entity in self.layout.entities}
# One batched PhysX contact view per requested entity reports the net
# contact force on every body against any contact object. Explicit
# per-body patterns are required because the baked USD nests bodies by
# kinematic depth, which a single-level sensor leaf pattern cannot
# match.
if self.net_contact_entities:
from isaacsim.core.simulation_manager import SimulationManager
physics_view = SimulationManager.get_physics_sim_view()
for entity_name in self.net_contact_entities:
entity = layout_entities[entity_name]
component = self.entity_components[entity_name]
body_paths = self.entity_body_paths[entity_indexes[entity_name]][0]
patterns = [
"/World/envs/env_*/" + component + body_paths[body_name]
for body_name in entity.body_names
]
body_view = physics_view.create_rigid_body_view(patterns)
contact_view = physics_view.create_rigid_contact_view(patterns)
count = self.num_envs * len(entity.body_names)
prim_paths = list(body_view.prim_paths)
if len(prim_paths) != count:
raise RuntimeError(
f"body-net contact reporter for entity {entity_name!r} resolved "
f"{len(prim_paths)} bodies; expected {count}"
)
relative_to_body = {
body_paths[name]: index for index, name in enumerate(entity.body_names)
}
env_rows = np.empty(count, dtype=np.int64)
body_columns = np.empty(count, dtype=np.int64)
seen: set[tuple[int, int]] = set()
positions = _environment_positions(self.entity_paths[entity_name])
for row, path in enumerate(prim_paths):
try:
env = _environment_index(path, self.entity_paths[entity_name], positions)
except RuntimeError:
raise RuntimeError(
"body-net contact reporter contains an unowned or ambiguous "
f"instance: {path}"
) from None
# entity_paths already include the entity component scope.
relative = path[len(self.entity_paths[entity_name][env]) :]
body = relative_to_body.get(relative)
if body is None or (env, body) in seen:
raise RuntimeError(
f"body-net contact reporter resolved an unexpected prim: {path}"
)
seen.add((env, body))
env_rows[row] = env
body_columns[row] = entity.body_ids[body]
self.net_contact_views.append(contact_view)
self.net_contact_maps.append(
{
"env": env_rows,
"body": body_columns,
"entity": entity_name,
"count": count,
# Keep the rigid-body view alive alongside the contact view.
"body_view": body_view,
}
)
telemetry.mark("sim_reset")
for entity, entry, asset in zip(self.layout.entities, self.entries, self.assets):
native_paths = list(asset.root_physx_view.prim_paths)
native_envs = _native_environment_order(native_paths, self.entity_paths[entity.name])
env_map = np.argsort(native_envs)
native_bodies = list(asset.body_names)
bodies = self.renderer._build_permutation(
native_bodies, list(entity.body_names), "body"
)
native_joints = list(asset.joint_names) if entity.kind == "articulation" else []
if entity.kind == "articulation" and (
bool(asset.is_fixed_base) != (entity.root_mode == "fixed")
):
raise RuntimeError("native articulation root mode differs from declaration")
joints = self.renderer._build_permutation(
native_joints, [joint.name for joint in entity.joints], "joint"
)
control_joints = np.asarray(
[native_joints.index(name) for name in entity.actuator_joint_names], dtype=np.int64
)
geom_columns = _native_geometry_columns(
native_bodies, bodies, entity, _consistent_native_mask(entry)
)
self.maps.append(
{
"bodies": bodies,
"geoms": geom_columns,
"joints": joints,
"controls": control_joints,
"envs": env_map,
"public_for_native": np.asarray(native_envs),
}
)
self._apply_variant_drives(entity, entry, asset, self.maps[-1])
telemetry.mark("maps_build")
for entity, entry, asset, mapping in zip(
self.layout.entities, self.entries, self.assets, self.maps
):
self._verify_joint_velocity_limits(entity, entry, asset, mapping)
telemetry.mark("joint_velocity_limits")
ids = np.arange(self.num_envs, dtype=np.int64)
self._commit(
ids,
np.asarray(payload["initial_qpos"], dtype=np.float32),
np.asarray(payload["initial_qvel"], dtype=np.float32),
np.asarray(payload["initial_roots"], dtype=np.float32),
np.ones(self.layout.nq, dtype=np.uint8),
np.ones(self.layout.nv, dtype=np.uint8),
np.ones((len(self.assets), 2), dtype=np.uint8),
initializing=True,
)
if "initial_ctrl" in payload:
self._set_control_targets(np.asarray(payload["initial_ctrl"], dtype=np.float32))
telemetry.mark("commit")
self.actual = self._audit_instances()
telemetry.mark("audit")
progress.report("isaacsim worker: finalizing", 1, 1, force=True)
telemetry.count_prims(self.sim.stage, "final")
try:
free, total = torch.cuda.mem_get_info(self.device)
telemetry.vram_used_mb = (total - free) / 1e6
except Exception: # telemetry must never fail INIT
telemetry.vram_used_mb = None
self._init_telemetry = telemetry.to_wire()
return self.get_meta()
[docs]
def adopt_initialized_context(
self, renderer: Any, metadata: dict[str, Any], payload: dict[str, Any]
) -> None:
"""Project the model-file wire onto this runtime after its existing cold importer.
This adopts live views and the already initialized Kit context. The
synthetic execution layout preserves the historical D-wide control
surface; it is never published as a new physical entity capability.
"""
bridge = self.protocol.load_model_file_projection()
self.renderer = renderer
self.num_envs, self.sim_dt = renderer.num_envs, renderer.sim_dt
self.sim, self.torch, self.device = renderer.sim, renderer.torch, renderer.device
self.origins = np.asarray(renderer.env_origins, dtype=np.float32).copy()
if self.origins.shape != (self.num_envs, 3) or not np.isfinite(self.origins).all():
raise RuntimeError("model-file native environment origins are malformed")
self.layout = bridge.ModelFileExecutionLayout(
renderer.contract_joint_names,
renderer.contract_body_names,
root_body_name=payload.get("root_body_name"),
)
self.model_file_projection = bridge.ModelFileSlotProjection(
self.protocol, self.num_envs, self.layout
)
self.assets = [renderer.robot]
# Resolve instance rows from the existing view rather than assuming
# GridCloner creation order equals PhysX view order.
roots = [path + "/Robot" for path in renderer.env_prim_paths]
native_envs = _native_environment_order(
list(renderer.robot.root_physx_view.prim_paths), roots
)
joints = np.asarray(renderer.native_joint_for_contract, dtype=np.int64).copy()
bodies = np.asarray(renderer.native_body_for_contract, dtype=np.int64).copy()
self.maps = [
{
"envs": np.argsort(native_envs),
"public_for_native": native_envs,
"joints": joints,
"bodies": bodies,
"controls": joints.copy(),
}
]
self._model_file_metadata = metadata.copy()
def _setup_renderer(self, sim_utils: Any, payload: dict[str, Any]) -> None:
owner = self.renderer
owner.sim, owner.torch, owner.device = self.sim, self.torch, self.device
owner.num_envs, owner.sim_dt = self.num_envs, self.sim_dt
owner.robot = self.assets[0]
owner.render_mode = payload.get("render_mode", "none")
owner.render_width = payload.get("render_width", 1280)
owner.render_height = payload.get("render_height", 720)
if owner.render_mode != "none":
light = sim_utils.DomeLightCfg(intensity=100.0)
light.func("/World/UniSimLight", light)
if owner.render_mode == "record":
from isaaclab.sensors.camera import Camera, CameraCfg
owner.camera = Camera(
CameraCfg(
prim_path="/World/envs/env_0/UniSimCamera",
update_period=0.0,
data_types=["rgb"],
width=owner.render_width,
height=owner.render_height,
spawn=sim_utils.PinholeCameraCfg(clipping_range=(0.1, 1e5)),
)
)
def _bind_fixed_anchors(self) -> None:
"""Place imported world joints at the actual cloned root world transform."""
from pxr import Gf, Usd, UsdGeom, UsdPhysics
for entity, asset in zip(self.layout.entities, self.assets):
if entity.kind != "articulation" or entity.root_mode != "fixed":
continue
for path in self.entity_paths[entity.name]:
prim = asset.stage.GetPrimAtPath(path)
fixed = []
for child in Usd.PrimRange(prim):
if not child.IsA(UsdPhysics.FixedJoint):
continue
joint = UsdPhysics.FixedJoint(child)
body0 = joint.GetBody0Rel().GetTargets()
body1 = joint.GetBody1Rel().GetTargets()
if not body0 and len(body1) == 1:
fixed.append((joint, body1[0]))
if len(fixed) != 1:
raise RuntimeError("fixed entity requires one world anchor joint")
joint, root_path = fixed[0]
transform = UsdGeom.XformCache().GetLocalToWorldTransform(
asset.stage.GetPrimAtPath(root_path)
)
position = transform.ExtractTranslation()
quaternion = transform.ExtractRotationQuat()
joint.CreateLocalPos0Attr().Set(Gf.Vec3f(*position))
joint.CreateLocalRot0Attr().Set(Gf.Quatf(quaternion))
joint.CreateLocalPos1Attr().Set(Gf.Vec3f(0))
joint.CreateLocalRot1Attr().Set(Gf.Quatf(1))
def _audit_instances(self) -> list[dict[str, Any]]:
"""Read identity and native properties from actual spawned instances."""
from pxr import PhysxSchema, Usd, UsdPhysics
result = []
for entity_index, (entity, entry, asset, mapping) in enumerate(
zip(self.layout.entities, self.entries, self.assets, self.maps)
):
paths = self.entity_paths[entity.name]
# prim_paths re-resolves the view pattern over the whole stage on
# every access (C++ getter); hoist it out of the per-row indexing.
view_prim_paths = asset.root_physx_view.prim_paths
native_paths = [view_prim_paths[i] for i in mapping["envs"]]
if not np.array_equal(
_native_environment_order(native_paths, paths), np.arange(self.num_envs)
):
raise RuntimeError("native view row order differs from environment order")
observed = []
for path in paths:
prim = asset.stage.GetPrimAtPath(path)
value = prim.GetAttribute("unisim:variantIndex").Get()
if not isinstance(value, int):
raise RuntimeError("spawned asset has no observable variant identity")
observed.append(value)
if observed != entry["assignment"]:
raise RuntimeError("actual spawned variant assignment differs from requested")
representatives = _assignment_representatives(observed)
masses = _numpy(asset.root_physx_view.get_masses()).reshape(self.num_envs, -1)[
mapping["envs"]
]
masses = masses[:, mapping["bodies"]]
expected = np.asarray([entry["variants"][index]["body_mass"] for index in observed])
if not np.allclose(masses, expected, rtol=2e-4, atol=1e-6):
raise RuntimeError(
f"entity {entity.name} native body masses differ: "
f"actual={masses.tolist()}, requested={expected.tolist()}"
)
coms = _numpy(asset.root_physx_view.get_coms()).reshape(self.num_envs, -1, 7)[
mapping["envs"]
]
coms = coms[:, mapping["bodies"]]
expected_coms = np.asarray(
[entry["variants"][variant]["body_ipos"] for variant in observed]
)
if not np.allclose(coms[:, :, :3], expected_coms, rtol=1e-4, atol=1e-6):
raise RuntimeError(f"entity {entity.name} native COM differs from source")
inertias = _numpy(asset.root_physx_view.get_inertias()).reshape(
self.num_envs, -1, 3, 3
)[mapping["envs"]][:, mapping["bodies"]]
expected_inertias = []
for variant in observed:
record = entry["variants"][variant]
matrices = []
for diagonal, quaternion in zip(record["body_inertia"], record["body_iquat"]):
# Columns of R are independently rotated unit basis vectors.
rotation = _rotate(np.broadcast_to(quaternion, (3, 4)), np.eye(3)).T
matrices.append((rotation * np.asarray(diagonal)) @ rotation.T)
expected_inertias.append(matrices)
if not np.allclose(inertias, expected_inertias, rtol=2e-4, atol=1e-6):
raise RuntimeError(f"entity {entity.name} native inertia differs from source")
variant_body_paths = self.entity_body_paths[entity_index][0]
def audit_destination(env_index: int) -> tuple[list[list[float]], dict[str, Any]]:
"""Full USD subtree audit of one destination environment."""
variant = observed[env_index]
path = paths[env_index]
prim = asset.stage.GetPrimAtPath(path)
if not entry["collision_enabled"]:
for child in Usd.PrimRange(prim):
if child.HasAPI(UsdPhysics.CollisionAPI) and (
UsdPhysics.CollisionAPI(child).GetCollisionEnabledAttr().Get()
):
raise RuntimeError("collision-disabled entity has an enabled collider")
if entity.kind == "articulation":
# The bake guarantees exactly one UsdPhysics articulation
# root per spawned clone: the importer's synthetic worldBody
# keeps an inert PhysxArticulationAPI after its
# ArticulationRootAPI is removed, so the flag must be read
# from the actual articulation root prim only. Environment
# collision filtering must not have clobbered it.
roots = [
child
for child in Usd.PrimRange(prim)
if child.HasAPI(UsdPhysics.ArticulationRootAPI)
]
if len(roots) != 1:
raise RuntimeError(
f"entity {entity.name} has ambiguous native articulation roots"
)
flag = (
PhysxSchema.PhysxArticulationAPI(roots[0])
.GetEnabledSelfCollisionsAttr()
.Get()
)
if bool(flag) != bool(entry["self_collision"]):
raise RuntimeError(
f"entity {entity.name} native self-collision differs from request"
)
row: list[list[float]] = []
for body_name in entity.body_names:
body_prim = asset.stage.GetPrimAtPath(path + variant_body_paths[body_name])
if not body_prim or not body_prim.IsValid():
raise RuntimeError(
f"entity {entity.name} native body prim is missing: {body_name}"
)
row.append(self._native_source_sphere_radii(body_prim, record, body_name))
expected_radii = entry["variants"][variant]["body_sphere_radii"]
if not body_sphere_radii_close(row, expected_radii, rtol=2e-6, atol=1e-8):
raise RuntimeError(
f"entity {entity.name} native sphere radii differ: "
f"actual={row}, requested={expected_radii}"
)
return row, _native_geometry_record(
asset.stage.GetPrimAtPath(path),
entity,
entry["variants"][variant],
variant_body_paths,
)
# Destinations of one variant are byte-identical Sdf.CopySpec
# copies of the same normalized prototype spec, and no code path
# authors per-environment USD divergence after the copy, so the
# structural subtree audit runs on one representative environment
# per assigned variant. The PhysX view readbacks (masses, COMs,
# inertias, drive gains, joint types, metatype topology) above and
# below still cover every environment. A representative failure
# falls back to a full per-environment scan so the error
# pinpoints every offending instance.
try:
measured = {
variant: audit_destination(env_index)
for variant, env_index in sorted(representatives.items())
}
except RuntimeError:
for env_index in range(self.num_envs):
audit_destination(env_index)
raise
sphere_radii = [measured[variant][0] for variant in observed]
# Geometry placement is environment-specific even when Sdf.CopySpec
# makes the collider structure identical. Audit every row so path,
# world transform, and world bounds are never copied from a
# representative environment; sphere-radius structure remains
# representative because it contains no placement data.
geometry_rows = [audit_destination(env_index)[1] for env_index in range(self.num_envs)]
if entity.joints:
kinds = _numpy(asset.root_physx_view.get_dof_types())[mapping["envs"]][
:, mapping["joints"]
]
expected_kinds = [0 if joint.kind == "hinge" else 1 for joint in entity.joints]
if not np.all(kinds == expected_kinds):
raise RuntimeError(f"entity {entity.name} native joint types differ")
for field, actual in (
("dof_stiffness", asset.root_physx_view.get_dof_stiffnesses()),
("dof_damping", asset.root_physx_view.get_dof_dampings()),
):
values = _numpy(actual)[mapping["envs"]][:, mapping["joints"]]
expected = np.asarray(
[entry["variants"][variant][field] for variant in observed]
)
if not np.allclose(values, expected, rtol=1e-5, atol=1e-6):
raise RuntimeError(f"entity {entity.name} native {field} differs")
for env in range(self.num_envs):
metatype = asset.root_physx_view.get_metatype(int(mapping["envs"][env]))
parents = dict(zip(metatype.link_names, metatype.link_parents))
for body, parent in zip(entity.body_names, entity.body_parent_names):
if parent is not None and parents[body] != parent:
raise RuntimeError(f"entity {entity.name} native body topology differs")
result.append(
{
"name": entity.name,
"assignment": observed,
"prim_paths": native_paths,
"body_mass": masses.tolist(),
"body_com": coms[:, :, :3].tolist(),
"body_inertia": inertias.tolist(),
"body_sphere_radii": sphere_radii,
"geom_names": [row["geom_names"] for row in geometry_rows],
"geom_body_names": [row["geom_body_names"] for row in geometry_rows],
"geom_contact_masks": [row["geom_contact_masks"] for row in geometry_rows],
"geom_friction": [row["geom_friction"] for row in geometry_rows],
"geometry_audit_schema_version": geometry_rows[0][
"geometry_audit_schema_version"
]
if geometry_rows
else 1,
"geom_source_types": [row["geom_source_types"] for row in geometry_rows],
"geom_source_sizes": [row["geom_source_sizes"] for row in geometry_rows],
"geom_source_poses": [row["geom_source_poses"] for row in geometry_rows],
"geom_native_audits": [row["geom_native_audits"] for row in geometry_rows],
}
)
return result
[docs]
def attach_slots(self, payload: dict[str, Any]) -> None:
from multiprocessing import resource_tracker, shared_memory
shapes = (
self.protocol.slot_shapes(
self.num_envs, len(self.layout.entities[0].joints), self.layout.nbody
)
if self.model_file_projection is not None
else self.protocol.scene_slot_shapes(
self.num_envs, self.layout, len(self.contact_force_sensors)
)
)
self.protocol.validate_slot_specs(payload["slots"], shapes)
attached = {}
for name, spec in payload["slots"].items():
handle = shared_memory.SharedMemory(name=spec["shm"], create=False)
resource_tracker.unregister(handle._name, "shared_memory") # type: ignore[attr-defined]
self._shm_handles.append(handle)
attached[name] = np.ndarray(
tuple(spec["shape"]), dtype=spec["dtype"], buffer=handle.buf
)
if self.model_file_projection is None:
self.slots = attached
else:
self.slots = self.model_file_projection.attach(attached)
self.refresh_state_slots()
[docs]
def attach_cuda_ipc(self, payload: dict[str, Any]) -> dict[str, Any]:
"""Import the host-owned canonical CUDA arena on a cold boundary."""
if self.model_file_projection is not None:
raise NotImplementedError("IsaacSim CUDA IPC requires a mapped scene")
if self._contact_reporting:
raise NotImplementedError("IsaacSim CUDA IPC does not support tensor sensors yet")
worker_device = self.torch.device(self.device)
if worker_device.type != "cuda":
raise RuntimeError(
f"IsaacSim CUDA IPC requires a CUDA worker, got device {worker_device}"
)
device_index = (
worker_device.index
if worker_device.index is not None
else int(self.torch.cuda.current_device())
)
arena = WorkerCudaIpcArena(payload, device_index=device_index)
try:
expected = arena.layout.shapes
for name in arena.layout.shapes:
view = getattr(arena, name)
require_cuda_tensor_of_dtype(
view,
rank=len(expected[name]),
name=name,
dtype=arena.layout.dtypes[name],
)
if tuple(view.shape) != expected[name]:
raise RuntimeError(f"IsaacSim CUDA arena {name} shape mismatch")
self._cuda_reset_sequence = 0
self._cuda_maps = []
for entity, asset, mapping in zip(self.layout.entities, self.assets, self.maps):
self._cuda_maps.append(
{
"asset": asset,
"rows": self.torch.arange(
self.num_envs, dtype=self.torch.long, device=self.device
),
"native_rows": self.torch.as_tensor(
mapping["envs"], dtype=self.torch.long, device=self.device
),
"public_for_native": self.torch.as_tensor(
mapping["public_for_native"],
dtype=self.torch.long,
device=self.device,
),
"joints": self.torch.as_tensor(
mapping["joints"], dtype=self.torch.long, device=self.device
),
"joint_ids": list(mapping["joints"]),
"public_body_ids": self.torch.as_tensor(
entity.body_ids,
dtype=self.torch.long,
device=self.device,
),
"bodies": self.torch.as_tensor(
mapping["bodies"],
dtype=self.torch.long,
device=self.device,
),
"root_qpos_columns": self.torch.as_tensor(
entity.root_qpos_indices,
dtype=self.torch.long,
device=self.device,
),
"root_qvel_columns": self.torch.as_tensor(
entity.root_qvel_indices,
dtype=self.torch.long,
device=self.device,
),
"joint_qpos_columns": self.torch.as_tensor(
[joint.qpos_indices[0] for joint in entity.joints],
dtype=self.torch.long,
device=self.device,
),
"joint_qvel_columns": self.torch.as_tensor(
[joint.qvel_indices[0] for joint in entity.joints],
dtype=self.torch.long,
device=self.device,
),
"controls": list(mapping["controls"]),
}
)
self._cuda_sensor_specs = self._bind_cuda_sensor_specs(payload["sensors"])
self._cuda_origins = self.torch.as_tensor(
self.origins, dtype=self.torch.float32, device=self.device
)
self._cuda_ipc = arena
arena.body_state.zero_()
arena.body_state[..., 3] = 1.0
self._publish_cuda_state()
except Exception:
view = None
arena.close()
self._cuda_ipc = None
self._cuda_maps = []
self._cuda_sensor_specs = {}
raise
return {"device_uuid": arena.device_uuid, "layout": arena.layout.to_wire()}
def _bind_cuda_sensor_specs(
self, sensor_specs: list[dict[str, Any]]
) -> dict[str, dict[str, Any]]:
if not isinstance(sensor_specs, list) or len(sensor_specs) > 2:
raise ValueError("IsaacSim CUDA IPC supports at most two scalar sensor projections")
allowed_names = {"pelvis_local_linvel", "torso_gyro"}
allowed_kinds = {"local_linvel", "gyro"}
bound: dict[str, dict[str, Any]] = {}
for spec in sensor_specs:
if not isinstance(spec, dict) or set(spec) != {
"name",
"kind",
"body_id",
"local_pos",
"local_quat",
}:
raise ValueError("malformed IsaacSim CUDA IPC sensor descriptor")
name = spec["name"]
kind = spec["kind"]
body_id = spec["body_id"]
if (
not isinstance(name, str)
or name not in allowed_names
or kind not in allowed_kinds
or isinstance(body_id, bool)
or not isinstance(body_id, int)
or body_id < 0
or body_id >= self.layout.nbody
or name in bound
):
raise ValueError("unsupported IsaacSim CUDA IPC sensor descriptor")
bound[name] = {
"kind": kind,
"body_id": body_id,
"local_pos": self.torch.as_tensor(
spec["local_pos"], dtype=self.torch.float32, device=self.device
),
"local_quat": self.torch.as_tensor(
spec["local_quat"], dtype=self.torch.float32, device=self.device
),
}
return bound
[docs]
def refresh_state_slots(self) -> None:
for field in ("qpos", "qvel", "entity_root_state", "body_state", "contact_force"):
self.slots[field].fill(0)
if "contact_sensor_force" in self.slots:
self.slots["contact_sensor_force"].fill(0)
# Unowned engine-world bodies still have a valid identity orientation.
self.slots["body_state"][:, :, 3] = 1
for index, (entity, asset, mapping) in enumerate(
zip(self.layout.entities, self.assets, self.maps)
):
root = _numpy(asset.data.root_link_state_w)[mapping["envs"]].copy()
root[:, :3] -= self.origins
bodies = _numpy(asset.data.body_link_state_w)[mapping["envs"]].copy()
bodies[:, :, :3] -= self.origins[:, None, :]
self.slots["entity_root_state"][:, index] = root
self.slots["body_state"][:, entity.body_ids] = bodies[:, mapping["bodies"]]
if entity.root_mode == "floating":
self.slots["qpos"][:, entity.root_qpos_indices] = root[:, :7]
velocity = root[:, 7:13].copy()
velocity[:, 3:] = _rotate(root[:, 3:7], velocity[:, 3:], inverse=True)
self.slots["qvel"][:, entity.root_qvel_indices] = velocity
if entity.joints:
pos = _numpy(asset.data.joint_pos)[mapping["envs"]][:, mapping["joints"]]
vel = _numpy(asset.data.joint_vel)[mapping["envs"]][:, mapping["joints"]]
self.slots["qpos"][:, [j.qpos_indices[0] for j in entity.joints]] = pos
self.slots["qvel"][:, [j.qvel_indices[0] for j in entity.joints]] = vel
if self.model_file_projection is not None:
self.model_file_projection.publish()
if getattr(self, "_cuda_ipc", None) is not None:
self._publish_cuda_state()
def _native_device_tensor(self, value: Any, name: str) -> Any:
if not isinstance(value, self.torch.Tensor) or not bool(value.is_cuda):
raise RuntimeError(
f"IsaacSim CUDA IPC requires native {name} to be a CUDA tensor, got {value!r}"
)
return value
def _rotate_cuda(self, q: Any, v: Any, *, inverse: bool) -> Any:
quat = q.clone()
if inverse:
quat[..., 1:4] = -quat[..., 1:4]
vector = quat[..., 1:4]
cross = self.torch.cross(vector, v, dim=-1)
return v + quat[..., 0:1] * (2.0 * cross) + 2.0 * self.torch.cross(vector, cross, dim=-1)
def _publish_cuda_body_state(self, selected_rows: Any | None = None) -> None:
arena = self._cuda_ipc
if arena is None:
return
for mapping in self._cuda_maps:
if "asset" not in mapping:
continue
body_ids = mapping["public_body_ids"]
if not int(body_ids.numel()):
continue
native_rows = mapping["native_rows"]
origins = self._cuda_origins
if selected_rows is not None:
native_rows = native_rows.index_select(0, selected_rows)
origins = origins.index_select(0, selected_rows)
bodies = self._native_device_tensor(
mapping["asset"].data.body_link_state_w, "body link state"
).index_select(0, native_rows)
bodies = bodies[:, mapping["bodies"]].clone()
bodies[..., 0:3].sub_(origins[:, None, :])
if selected_rows is None:
arena.body_state.index_copy_(1, body_ids, bodies)
else:
arena.body_state[selected_rows[:, None], body_ids] = bodies
def _publish_cuda_scalar_sensors(self, selected_rows: Any | None = None) -> None:
arena = self._cuda_ipc
if arena is None or not getattr(self, "_cuda_sensor_specs", {}):
return
for name, spec in self._cuda_sensor_specs.items():
body = (
arena.body_state[:, spec["body_id"]]
if selected_rows is None
else arena.body_state[selected_rows, spec["body_id"]]
)
if spec["kind"] == "local_linvel":
offset_world = self._rotate_cuda(
body[:, 3:7],
spec["local_pos"].unsqueeze(0).expand(body.shape[0], 3),
inverse=False,
)
vector = body[:, 7:10] + self.torch.cross(body[:, 10:13], offset_world, dim=-1)
else:
vector = body[:, 10:13]
body_frame = self._rotate_cuda(body[:, 3:7], vector, inverse=True)
slot = 0 if name == "pelvis_local_linvel" else 1
local_quat = spec["local_quat"].unsqueeze(0).expand(body_frame.shape[0], 4)
if selected_rows is None:
arena.sensor_state[:, slot] = self._rotate_cuda(
local_quat,
body_frame,
inverse=True,
)
else:
arena.sensor_state[selected_rows, slot] = self._rotate_cuda(
local_quat,
body_frame,
inverse=True,
)
def _set_control_tensor_targets(self, control: Any) -> None:
"""Apply public control columns without converting them to NumPy."""
require_cuda_tensor(control, rank=2, name="control")
if tuple(control.shape) != (self.num_envs, self.layout.nu):
raise ValueError(
f"control must have shape {(self.num_envs, self.layout.nu)}, "
f"got {tuple(control.shape)}"
)
for entity, asset, mapping in zip(self.layout.entities, self.assets, self._cuda_maps):
if entity.actuator_indices:
native_control = control.index_select(0, mapping["public_for_native"])
asset.set_joint_position_target(
native_control[:, entity.actuator_indices].contiguous(),
joint_ids=mapping["controls"],
)
def _publish_cuda_state(self, selected_rows: Any | None = None) -> None:
"""Project native IsaacLab state directly into canonical CUDA views."""
arena = self._cuda_ipc
if arena is None:
return
for entity, asset, mapping in zip(self.layout.entities, self.assets, self._cuda_maps):
rows = mapping["rows"] if selected_rows is None else selected_rows
native_rows = mapping["native_rows"]
if selected_rows is not None:
native_rows = native_rows.index_select(0, selected_rows)
if entity.root_mode == "floating":
root = self._native_device_tensor(
asset.data.root_link_state_w, "root link state"
).index_select(0, native_rows)
root = root.clone()
root[:, 0:3].sub_(self._cuda_origins.index_select(0, rows))
if selected_rows is None:
arena.qpos.index_copy_(1, mapping["root_qpos_columns"], root[:, 0:7])
else:
arena.qpos[selected_rows[:, None], mapping["root_qpos_columns"]] = root[:, 0:7]
velocity = root[:, 7:13].clone()
velocity[:, 3:6] = self._rotate_cuda(root[:, 3:7], velocity[:, 3:6], inverse=True)
if selected_rows is None:
arena.qvel.index_copy_(1, mapping["root_qvel_columns"], velocity)
else:
arena.qvel[selected_rows[:, None], mapping["root_qvel_columns"]] = velocity
if entity.joints:
positions = self._native_device_tensor(
asset.data.joint_pos, "joint positions"
).index_select(0, native_rows)
velocities = self._native_device_tensor(
asset.data.joint_vel, "joint velocities"
).index_select(0, native_rows)
if mapping["joints"].numel():
positions = positions.index_select(1, mapping["joints"])
velocities = velocities.index_select(1, mapping["joints"])
if selected_rows is None:
arena.qpos.index_copy_(1, mapping["joint_qpos_columns"], positions)
arena.qvel.index_copy_(1, mapping["joint_qvel_columns"], velocities)
else:
arena.qpos[selected_rows[:, None], mapping["joint_qpos_columns"]] = positions
arena.qvel[selected_rows[:, None], mapping["joint_qvel_columns"]] = velocities
self._publish_cuda_body_state(selected_rows)
self._publish_cuda_scalar_sensors(selected_rows)
arena.record_state()
# The pipe READY reply is a worker-health/lifecycle acknowledgement,
# not a data barrier. This event carries the asynchronous D2D
# projection order to the host consumer stream.
def _publish_cuda_reset_state(self, rows: Any, qpos: Any, qvel: Any) -> None:
"""Publish canonical reset state without a native qpos/qvel projection."""
arena = self._cuda_ipc
if arena is None:
raise RuntimeError("IsaacSim CUDA IPC arena is not attached")
# CompiledSceneLayout owns the no-gap/no-overlap qpos/qvel invariant,
# so reset rows can cross this boundary as whole canonical rows.
arena.qpos.index_copy_(0, rows, qpos)
arena.qvel.index_copy_(0, rows, qvel)
self._publish_cuda_body_state(rows)
self._publish_cuda_scalar_sensors(rows)
arena.record_state()
def _apply_cuda_reset_properties(self, rows: Any, count: int, terms: tuple[str, ...]) -> None:
"""Cross the explicit reset-only PhysX property bridge."""
arena = self._cuda_ipc
if arena is None:
raise RuntimeError("IsaacSim CUDA IPC arena is not attached")
raw: dict[str, Any] = {}
if "kp" in terms:
raw["kp"] = arena.reset_kp[:count].detach().cpu().numpy()
if "kd" in terms:
raw["kd"] = arena.reset_kd[:count].detach().cpu().numpy()
randomization = self._validated_reset_randomization({"randomization": raw}, count)
if randomization is None:
return
host_rows = rows.detach().cpu().numpy()
before_properties = self._readback_reset_properties()
self._apply_reset_randomization(host_rows, randomization)
property_records = self._readback_reset_properties()
self._verify_reset_property_readback(
host_rows,
randomization,
before_properties,
property_records,
)
[docs]
def reset_cuda_ipc(self, payload: dict[str, Any]) -> dict[str, Any]:
"""Apply a selected canonical CUDA reset and republish device state."""
if self.faulted:
raise RuntimeError("IsaacSim CUDA IPC scene is faulted")
arena = self._cuda_ipc
if arena is None:
raise RuntimeError("IsaacSim CUDA IPC arena is not attached")
if not isinstance(payload, dict) or not {"count", "sequence"} <= set(payload):
raise ValueError("IsaacSim CUDA IPC reset payload must contain count/sequence")
if set(payload) - {"count", "sequence", "has_randomization", "randomization_terms"}:
raise ValueError("IsaacSim CUDA IPC reset payload contains unsupported fields")
count = payload["count"]
sequence = payload["sequence"]
if (
isinstance(count, bool)
or not isinstance(count, int)
or count <= 0
or count > self.num_envs
):
raise ValueError("invalid IsaacSim CUDA IPC reset count")
if (
isinstance(sequence, bool)
or not isinstance(sequence, int)
or sequence != self._cuda_reset_sequence + 1
):
raise ValueError("invalid IsaacSim CUDA IPC reset sequence")
started = time.perf_counter()
arena.wait_reset()
rows = arena.reset_env_indices[:count].long()
qpos = arena.reset_qpos[:count]
qvel = arena.reset_qvel[:count]
try:
for entity, asset, mapping in zip(self.layout.entities, self.assets, self._cuda_maps):
if entity.root_mode != "floating" and not entity.joints:
continue
native_rows = mapping["native_rows"].index_select(0, rows)
if entity.root_mode == "floating":
pose = qpos.index_select(1, mapping["root_qpos_columns"]).contiguous()
pose[:, 0:3].add_(self._cuda_origins.index_select(0, rows))
asset.write_root_pose_to_sim(pose, env_ids=native_rows)
velocity = qvel.index_select(1, mapping["root_qvel_columns"]).contiguous()
velocity = velocity.clone()
velocity[:, 3:6] = self._rotate_cuda(
pose[:, 3:7], velocity[:, 3:6], inverse=False
)
asset.write_root_link_velocity_to_sim(velocity, env_ids=native_rows)
if entity.joints:
joint_ids = mapping["joint_ids"]
positions = qpos.index_select(1, mapping["joint_qpos_columns"]).contiguous()
velocities = qvel.index_select(1, mapping["joint_qvel_columns"]).contiguous()
asset.write_joint_state_to_sim(
positions,
velocities,
joint_ids=joint_ids,
env_ids=native_rows,
)
asset.reset(native_rows)
for asset in self.assets:
asset.update(self.sim_dt)
self._cuda_reset_sequence = sequence
property_started = time.perf_counter()
if payload.get("has_randomization"):
self._apply_cuda_reset_properties(rows, count, payload["randomization_terms"])
property_ms = (time.perf_counter() - property_started) * 1000.0
self._publish_cuda_reset_state(rows, qpos, qvel)
except Exception:
self.faulted = True
raise
return {
"timing": {
"reset_upload_ms": (time.perf_counter() - started) * 1000.0,
"reset_property_ms": property_ms,
"state_publish_ms": 0.0,
"cuda_ipc": True,
}
}
[docs]
def step_cuda_ipc(self, payload: dict[str, Any]) -> dict[str, Any]:
"""Run the minimal qpos/qvel/control path with CUDA IPC bulk data."""
if self.faulted:
raise RuntimeError("IsaacSim CUDA IPC scene is faulted")
arena = self._cuda_ipc
if arena is None:
raise RuntimeError("IsaacSim CUDA IPC arena is not attached")
if self._contact_reporting:
raise NotImplementedError("IsaacSim CUDA IPC does not support tensor sensors yet")
count = payload.get("nsteps")
if isinstance(count, bool) or not isinstance(count, int) or count <= 0:
raise ValueError("nsteps must be a positive integer")
if "body_wrench" in payload:
raise NotImplementedError("IsaacSim CUDA IPC does not support body wrench yet")
started = time.perf_counter()
arena.wait_control()
try:
self._set_control_tensor_targets(arena.ctrl)
for _ in range(count):
for asset in self.assets:
asset.write_data_to_sim()
self.sim.step(render=False)
for asset in self.assets:
asset.update(self.sim_dt)
self._publish_cuda_state()
except Exception:
self.faulted = True
raise
return {
"timing": {
"physics_ms": (time.perf_counter() - started) * 1000.0,
"cuda_ipc": True,
}
}
def _refresh_contact_sensor_forces(self) -> None:
"""Publish final-substep filtered normal forces in world coordinates."""
if not self.contact_sensors:
return
if "contact_sensor_force" not in self.slots:
raise RuntimeError("contact sensors were initialized without their shared-memory slot")
for index, (sensor, mapping) in enumerate(
zip(self.contact_sensors, self.contact_sensor_maps)
):
matrix = _numpy(sensor.data.force_matrix_w)
if matrix.shape != (self.num_envs, 1, 1, 3):
raise RuntimeError(
f"contact sensor {index} returned shape {matrix.shape}; expected "
f"{(self.num_envs, 1, 1, 3)}"
)
forces = matrix.reshape(self.num_envs, 3)[mapping["envs"]]
if not np.isfinite(forces).all():
raise RuntimeError(f"contact sensor {index} returned non-finite force")
self.slots["contact_sensor_force"][:, index, :] = forces
def _poll_body_net_contact_forces(self) -> None:
"""Read the raw PhysX views on every physics substep.
On the PhysX GPU backend a body's net-contact entry is zeroed only on
the exact step where its contact is lost; a reader that skips that
step keeps the last in-contact force indefinitely (IsaacLab issue
7613). Polling every substep keeps the device buffer current; the
getter rewrites one cached on-device tensor, so this stays off the
host-copy path.
"""
for view in self.net_contact_views:
view.get_net_contact_forces(dt=self.sim_dt)
def _refresh_body_net_contact_forces(self) -> None:
"""Publish final-substep per-body net contact forces in world coordinates."""
if not self.net_contact_views:
return
for view, mapping in zip(self.net_contact_views, self.net_contact_maps):
net = _numpy(view.get_net_contact_forces(dt=self.sim_dt))
expected = (mapping["count"], 3)
if net.shape != expected:
raise RuntimeError(
f"body-net contact reporter for entity {mapping['entity']!r} returned "
f"shape {net.shape}; expected {expected}"
)
if not np.isfinite(net).all():
raise RuntimeError(
f"body-net contact reporter for entity {mapping['entity']!r} returned "
"non-finite force"
)
self.slots["contact_force"][mapping["env"], mapping["body"]] = net
[docs]
def set_state(self, payload: dict[str, Any]) -> dict[str, Any]:
"""Translate only the model-file wire; native reset always uses reset_entities."""
if self.model_file_projection is None:
raise NotImplementedError("mapped scenes require RESET_ENTITIES")
count = payload["count"]
if isinstance(count, bool) or not isinstance(count, int) or not 0 <= count <= self.num_envs:
raise ValueError("invalid reset count")
if count == 0:
return {"timing": {}}
return self.reset_entities(self.model_file_projection.prepare_reset(count))
[docs]
def step(self, payload: dict[str, Any]) -> dict[str, Any]:
count = payload["nsteps"]
if isinstance(count, bool) or not isinstance(count, int) or count <= 0:
raise ValueError("nsteps must be a positive integer")
if not np.isfinite(self.slots["ctrl"]).all():
raise ValueError("control must be finite")
wrench = None
if "body_wrench" in payload:
expected = (self.num_envs, self.layout.nbody, 6)
encoded = payload["body_wrench"]
expected_nbytes = int(np.prod(expected, dtype=np.int64)) * np.dtype(np.float32).itemsize
if not isinstance(encoded, bytes) or len(encoded) != expected_nbytes:
raise ValueError("body wrench payload must be C-order float32 bytes")
wrench = np.frombuffer(encoded, dtype=np.float32).reshape(expected)
if wrench.shape != expected:
raise ValueError(f"body wrench must have shape {expected}, got {wrench.shape}")
if not np.isfinite(wrench).all():
raise ValueError("body wrench contains NaN or Inf")
started = time.perf_counter()
try:
self._set_control_targets(self.slots["ctrl"])
if wrench is not None:
self._stage_body_wrench(wrench)
for _ in range(count):
for asset in self.assets:
asset.write_data_to_sim()
self.sim.step(render=False)
for asset in self.assets:
asset.update(self.sim_dt)
for sensor in self.contact_sensors:
sensor.update(self.sim_dt)
self._poll_body_net_contact_forces()
if wrench is not None:
self._clear_body_wrench()
self.refresh_state_slots()
self._refresh_contact_sensor_forces()
self._refresh_body_net_contact_forces()
except Exception:
self.faulted = True
if wrench is not None:
try:
self._clear_body_wrench()
except Exception:
pass
raise
return {"timing": {"physics_ms": (time.perf_counter() - started) * 1000}}
def _stage_body_wrench(self, wrench: np.ndarray) -> None:
"""Set one public world-frame wrench table on the native asset views."""
for entity, asset, mapping in zip(self.layout.entities, self.assets, self.maps):
native_values = wrench[mapping["public_for_native"]][:, entity.body_ids, :]
force = self._tensor(native_values[..., 0:3])
torque = self._tensor(native_values[..., 3:6])
asset.set_external_force_and_torque(
force,
torque,
body_ids=mapping["bodies"].tolist(),
is_global=True,
)
def _clear_body_wrench(self) -> None:
"""Disable external-wrench buffers after a completed or failed step."""
for asset in self.assets:
zero = self.torch.zeros(
(self.num_envs, asset.num_bodies, 3),
dtype=self.torch.float32,
device=self.device,
)
asset.set_external_force_and_torque(zero, zero, is_global=True)
def _set_control_targets(self, control: np.ndarray) -> None:
"""Apply actuator columns; keyframe controls are independent of joint positions."""
for entity, asset, mapping in zip(self.layout.entities, self.assets, self.maps):
if entity.actuator_indices:
asset.set_joint_position_target(
self._tensor(control[mapping["public_for_native"]][:, entity.actuator_indices]),
joint_ids=mapping["controls"].tolist(),
)
def _apply_variant_drives(
self, entity: Any, entry: dict[str, Any], asset: Any, mapping: dict[str, Any]
) -> None:
"""Write every environment's assigned variant drive gains into PhysX.
``ImplicitActuatorCfg`` seeds one shared gain table at construction;
per-variant stiffness/damping then land through the per-row view
setters so the post-init audit can compare native values against each
environment's assigned variant. The tracked drive/natural damping
split is the composition baseline for reset-time kp/kd/dof_damping DR.
"""
if entity.kind != "articulation" or not entity.joints:
self._drive_damping.append(None)
self._natural_damping.append(None)
return
assignment = _validated_assignment(entry, self.num_envs)
stiffness = np.asarray(
[entry["variants"][int(variant)]["dof_stiffness"] for variant in assignment],
dtype=np.float32,
)
damping = np.asarray(
[entry["variants"][int(variant)]["dof_damping"] for variant in assignment],
dtype=np.float32,
)
native_ids = self.torch.as_tensor(
mapping["envs"], dtype=self.torch.long, device=self.device
)
joint_ids = mapping["joints"].tolist()
asset.write_joint_stiffness_to_sim(
self._tensor(stiffness), joint_ids=joint_ids, env_ids=native_ids
)
asset.write_joint_damping_to_sim(
self._tensor(damping), joint_ids=joint_ids, env_ids=native_ids
)
self._drive_damping.append(damping)
self._natural_damping.append(np.zeros_like(damping))
def _commit(
self,
ids: np.ndarray,
qpos: np.ndarray,
qvel: np.ndarray,
roots: np.ndarray,
pmask: np.ndarray,
vmask: np.ndarray,
rmask: np.ndarray,
*,
initializing: bool = False,
) -> None:
try:
for index, (entity, asset, mapping) in enumerate(
zip(self.layout.entities, self.assets, self.maps)
):
touched = bool(rmask[index].any())
pcols = [j.qpos_indices[0] for j in entity.joints]
vcols = [j.qvel_indices[0] for j in entity.joints]
selected = np.flatnonzero(pmask[pcols] | vmask[vcols])
if not touched and not selected.size:
continue
native_rows = mapping["envs"][ids]
native_ids = self.torch.as_tensor(
native_rows, dtype=self.torch.long, device=self.device
)
if rmask[index, 0] and (
entity.root_mode != "fixed" or initializing and entity.kind == "rigid"
):
pose = roots[:, index, :7].copy()
pose[:, :3] += self.origins[ids]
asset.write_root_pose_to_sim(self._tensor(pose), env_ids=native_ids)
if rmask[index, 1] and entity.root_mode == "floating":
asset.write_root_link_velocity_to_sim(
self._tensor(roots[:, index, 7:]), env_ids=native_ids
)
if entity.joints:
if selected.size:
touched = True
joint_ids = mapping["joints"][selected].tolist()
native_joints = self.torch.as_tensor(
joint_ids, dtype=self.torch.long, device=self.device
)
# Gather on-device before crossing the CPU boundary;
# a sparse reset must not download every environment.
positions = _numpy(
asset.data.joint_pos[native_ids[:, None], native_joints]
).copy()
velocities = _numpy(
asset.data.joint_vel[native_ids[:, None], native_joints]
).copy()
for column, public_index in enumerate(selected):
if pmask[pcols[public_index]]:
positions[:, column] = qpos[:, pcols[public_index]]
if vmask[vcols[public_index]]:
velocities[:, column] = qvel[:, vcols[public_index]]
asset.write_joint_state_to_sim(
self._tensor(positions),
self._tensor(velocities),
joint_ids=joint_ids,
env_ids=native_ids,
)
if touched:
asset.reset(native_ids)
asset.update(self.sim_dt)
except Exception:
self.faulted = True
raise
def _validated_reset_randomization(
self, payload: dict[str, Any], count: int
) -> dict[str, np.ndarray] | None:
if "randomization" not in payload:
return None
raw = payload["randomization"]
allowed = {
"geom_friction",
"body_mass",
"body_ipos",
"body_inertia",
"kp",
"kd",
"dof_damping",
"dof_armature",
"dof_frictionloss",
}
if not isinstance(raw, dict) or not set(raw) <= allowed:
raise ValueError("randomization must contain only supported property terms")
if getattr(self, "_share_friction_materials", False) and "geom_friction" in raw:
raise ValueError(
"geom_friction reset is immutable while share_friction_materials is enabled"
)
def wire_float_table(value: object) -> np.ndarray:
if not isinstance(value, (list, np.ndarray)):
raise ValueError("randomization property tables must be lists or arrays")
try:
return np.asarray(value, dtype=np.float32)
except (TypeError, ValueError) as exc:
raise ValueError("randomization property tables must be numeric") from exc
owned_bodies = np.sort(
np.concatenate(
[
np.asarray(list(entity.body_ids), dtype=np.int64)
for entity in self.layout.entities
]
)
)
def body_table(term: str, width: int | None, *, positive: bool) -> np.ndarray:
values = wire_float_table(raw[term])
expected = (
(count, self.layout.nbody) if width is None else (count, self.layout.nbody, width)
)
if values.shape != expected or not np.isfinite(values).all():
raise ValueError(f"randomization {term} must be finite with shape {expected}")
if positive:
# Unowned public rows (for example the world body) carry zero
# canonical placeholders and are never written.
checked = values[:, owned_bodies] if width is None else values[:, owned_bodies, :]
if np.any(checked <= 0.0):
raise ValueError(f"randomization {term} must be strictly positive")
return np.asarray(values, dtype=np.float32).copy()
result: dict[str, np.ndarray] = {}
if "geom_friction" in raw:
values = wire_float_table(raw["geom_friction"])
expected = (count, self.layout.ngeom, 3)
if values.shape != expected or not np.isfinite(values).all() or np.any(values < 0.0):
raise ValueError(
"randomization geom_friction must be finite nonnegative (count, ngeom, 3)"
)
friction = np.asarray(values, dtype=np.float32).copy()
if np.any(friction[..., 0] != friction[..., 1]) or np.any(friction[..., 2] != 0.0):
raise ValueError(
"randomization geom_friction requires equal static/dynamic and zero torsion"
)
result["geom_friction"] = friction
if "body_mass" in raw:
result["body_mass"] = body_table("body_mass", None, positive=True)
if "body_ipos" in raw:
result["body_ipos"] = body_table("body_ipos", 3, positive=False)
if "body_inertia" in raw:
result["body_inertia"] = body_table("body_inertia", 3, positive=True)
for term in ("kp", "kd"):
if term in raw:
values = wire_float_table(raw[term])
gain_shape = (count, self.layout.nu)
if (
values.shape != gain_shape
or not np.isfinite(values).all()
or np.any(values < 0.0)
):
raise ValueError(
f"randomization {term} must be finite nonnegative with shape {gain_shape}"
)
result[term] = np.asarray(values, dtype=np.float32).copy()
root_columns = [
column for entity in self.layout.entities for column in entity.root_qvel_indices
]
for term in ("dof_damping", "dof_armature", "dof_frictionloss"):
if term in raw:
values = wire_float_table(raw[term])
dof_shape = (count, self.layout.nv)
if (
values.shape != dof_shape
or not np.isfinite(values).all()
or np.any(values < 0.0)
):
raise ValueError(
f"randomization {term} must be finite nonnegative with shape {dof_shape}"
)
if root_columns and np.any(values[:, root_columns] != 0.0):
raise ValueError(
f"randomization {term} free-root columns must remain zero; "
"PhysX exposes no root DOF damping/armature/friction"
)
result[term] = np.asarray(values, dtype=np.float32).copy()
return result or None
def _native_mass_rows(self, asset: Any, mapping: dict[str, Any]) -> np.ndarray:
masses = _numpy(asset.root_physx_view.get_masses()).reshape(self.num_envs, -1).copy()
if masses.shape[1] < len(mapping["bodies"]):
raise RuntimeError("native rigid-body view is narrower than the frozen entity map")
return masses
def _native_com_rows(self, asset: Any, mapping: dict[str, Any]) -> np.ndarray:
coms = _numpy(asset.root_physx_view.get_coms()).reshape(self.num_envs, -1, 7).copy()
if coms.shape[1] < len(mapping["bodies"]):
raise RuntimeError("native rigid-body COM view is narrower than the frozen entity map")
return coms
def _native_inertia_rows(self, asset: Any, mapping: dict[str, Any]) -> np.ndarray:
inertias = _numpy(asset.root_physx_view.get_inertias()).reshape(self.num_envs, -1, 9).copy()
if inertias.shape[1] < len(mapping["bodies"]):
raise RuntimeError(
"native rigid-body inertia view is narrower than the frozen entity map"
)
return inertias
def _native_material_rows(
self, asset: Any, mapping: dict[str, Any], geom_count: int
) -> np.ndarray:
if not geom_count:
return np.empty((self.num_envs, 0, 3), dtype=np.float32)
materials = (
_numpy(asset.root_physx_view.get_material_properties())
.reshape(self.num_envs, -1, 3)
.copy()
)
if materials.shape[1] != geom_count:
raise RuntimeError("native material view does not match the frozen geometry map")
columns = mapping["geoms"]
return materials[mapping["envs"]][:, columns[columns >= 0]]
@staticmethod
def _native_dof_friction(asset: Any) -> Any:
view = asset.root_physx_view
properties = getattr(view, "get_dof_friction_properties", None)
if properties is None:
return view.get_dof_friction_coefficients()
# Isaac Sim >= 5.0 writes joint friction through [static, dynamic,
# viscous] property triplets; reset DR only composes the static term.
return properties()[..., 0]
def _apply_reset_randomization(
self, ids: np.ndarray, randomization: dict[str, np.ndarray]
) -> None:
"""Write selected public property rows through public PhysX views."""
geom_offset = 0
for index, (entity, asset, mapping) in enumerate(
zip(self.layout.entities, self.assets, self.maps)
):
native_rows = mapping["envs"][ids]
native_ids = self.torch.as_tensor(native_rows, dtype=self.torch.int32, device="cpu")
if "body_mass" in randomization:
masses = self._native_mass_rows(asset, mapping)
masses[native_rows[:, None], mapping["bodies"]] = randomization["body_mass"][
:, list(entity.body_ids)
]
asset.root_physx_view.set_masses(self._cpu_tensor(masses), indices=native_ids)
if "body_ipos" in randomization:
coms = self._native_com_rows(asset, mapping)
coms[native_rows[:, None], mapping["bodies"], 0:3] = randomization["body_ipos"][
:, list(entity.body_ids), :
]
asset.root_physx_view.set_coms(self._cpu_tensor(coms), indices=native_ids)
if "body_inertia" in randomization:
blocks = self._native_inertia_rows(asset, mapping)
for row, env in enumerate(ids):
matrices = self._expected_inertia_blocks(
index, int(env), randomization["body_inertia"][row]
)
blocks[int(native_rows[row]), mapping["bodies"], :] = matrices.reshape(-1, 9)
asset.root_physx_view.set_inertias(self._cpu_tensor(blocks), indices=native_ids)
count = len(entity.geoms)
if "geom_friction" in randomization and count:
materials = (
_numpy(asset.root_physx_view.get_material_properties())
.reshape(self.num_envs, -1, 3)
.copy()
)
columns = mapping["geoms"]
# Visual-only geoms own no native shape; their DR rows are
# accepted in the public layout but have nowhere to land.
valid = columns >= 0
materials[native_rows[:, None], columns[valid]] = randomization["geom_friction"][
:, geom_offset : geom_offset + count
][:, valid]
asset.root_physx_view.set_material_properties(
self._cpu_tensor(materials), indices=native_ids
)
self._apply_dof_randomization(entity, asset, mapping, index, ids, randomization)
geom_offset += len(entity.geoms)
def _expected_inertia_blocks(
self, entity_index: int, env: int, public_diagonals: np.ndarray
) -> np.ndarray:
"""Compose symmetric COM-frame inertia tensors for one environment.
Public ``body_inertia`` rows are diagonal and carry no orientation; the
principal-axes rotation stays the environment's immutable assigned
variant ``body_iquat`` (reset-time ``body_iquat`` DR fails closed).
"""
entity = self.layout.entities[entity_index]
entry = self.entries[entity_index]
assignment = _validated_assignment(entry, self.num_envs)
record = entry["variants"][int(assignment[env])]
diagonals = np.asarray(public_diagonals, dtype=np.float64)[list(entity.body_ids)]
quaternions = np.asarray(record["body_iquat"], dtype=np.float64)
matrices = []
for diagonal, quaternion in zip(diagonals, quaternions):
# Columns of R are independently rotated unit basis vectors.
rotation = _rotate(np.broadcast_to(quaternion, (3, 4)), np.eye(3)).T
matrices.append((rotation * diagonal) @ rotation.T)
return np.asarray(matrices, dtype=np.float32)
def _apply_dof_randomization(
self,
entity: Any,
asset: Any,
mapping: dict[str, Any],
index: int,
ids: np.ndarray,
randomization: dict[str, np.ndarray],
) -> None:
"""Write drive/natural joint parameters for one articulation entity."""
if entity.kind != "articulation" or not entity.joints:
return
joint_columns = [joint.qvel_indices[0] for joint in entity.joints]
actuator_positions = [
position
for position, joint in enumerate(entity.joints)
if joint.name in entity.actuator_joint_names
]
native_ids = self.torch.as_tensor(
mapping["envs"][ids], dtype=self.torch.long, device=self.device
)
kp = randomization.get("kp")
if kp is not None and entity.actuator_indices:
asset.write_joint_stiffness_to_sim(
self._tensor(kp[:, list(entity.actuator_indices)]),
joint_ids=mapping["controls"].tolist(),
env_ids=native_ids,
)
kd = randomization.get("kd")
dof_damping = randomization.get("dof_damping")
if kd is not None or dof_damping is not None:
drive = self._drive_damping[index]
natural = self._natural_damping[index]
assert drive is not None and natural is not None
if kd is not None and actuator_positions:
drive[ids[:, None], actuator_positions] = kd[:, list(entity.actuator_indices)]
if dof_damping is not None:
natural[ids] = dof_damping[:, joint_columns]
asset.write_joint_damping_to_sim(
self._tensor(drive[ids] + natural[ids]),
joint_ids=mapping["joints"].tolist(),
env_ids=native_ids,
)
armature = randomization.get("dof_armature")
if armature is not None:
asset.write_joint_armature_to_sim(
self._tensor(armature[:, joint_columns]),
joint_ids=mapping["joints"].tolist(),
env_ids=native_ids,
)
friction = randomization.get("dof_frictionloss")
if friction is not None:
asset.write_joint_friction_coefficient_to_sim(
self._tensor(friction[:, joint_columns]),
joint_ids=mapping["joints"].tolist(),
env_ids=native_ids,
)
def _verify_joint_velocity_limits(
self, entity: Any, entry: dict[str, Any], asset: Any, mapping: dict[str, Any]
) -> None:
"""Fail closed unless the engine's velocity caps match the INIT request.
IsaacLab applies ``ImplicitActuatorCfg.velocity_limit_sim`` through the
PhysX view's max-velocity channel; read the same channel back instead
of trusting the configuration. The caps come from one shared actuator
table, so every environment row must equal the declared limit.
"""
record = entry["variants"][0]
if "dof_velocity_limit" not in record:
return
getter = getattr(asset.root_physx_view, "get_dof_max_velocities", None)
if getter is None:
raise RuntimeError(
f"entity {entity.name} requested joint velocity limits but the PhysX "
"view exposes no max-velocity readback channel"
)
table = _numpy(getter())[mapping["envs"]][:, mapping["joints"]]
expected = np.asarray(record["dof_velocity_limit"], dtype=np.float64)
if table.shape != (self.num_envs, expected.shape[0]) or not np.isfinite(table).all():
raise RuntimeError(
f"entity {entity.name} native dof_velocity_limit readback is invalid"
)
if not np.allclose(table, expected[None, :], rtol=2e-5, atol=1e-6):
raise RuntimeError(
f"entity {entity.name} native dof_velocity_limit readback differs from the "
f"INIT request: expected {expected.tolist()}"
)
def _readback_reset_properties(self) -> list[dict[str, Any]]:
records: list[dict[str, Any]] = []
for entity, asset, mapping in zip(self.layout.entities, self.assets, self.maps):
masses = self._native_mass_rows(asset, mapping)[mapping["envs"]][:, mapping["bodies"]]
coms = self._native_com_rows(asset, mapping)[mapping["envs"]][:, mapping["bodies"]]
inertias = self._native_inertia_rows(asset, mapping).reshape(self.num_envs, -1, 3, 3)[
mapping["envs"]
][:, mapping["bodies"]]
native_friction = self._native_material_rows(
asset, mapping, int((mapping["geoms"] >= 0).sum())
)
current = (
self.actual[entity.name]
if isinstance(self.actual, dict)
else next(item for item in self.actual if item["name"] == entity.name)
)
friction = (
np.asarray(current["geom_friction"], dtype=np.float32)
.reshape(self.num_envs, len(entity.geoms), 3)
.copy()
)
# PhysX owns values only for native collision geoms. Preserve the
# public visual-only slots so reset barriers continue to use the
# full-row contract used by the initial geometry audit.
friction[:, mapping["geoms"] >= 0, :] = native_friction
if (
not np.isfinite(masses).all()
or not np.isfinite(coms).all()
or not np.isfinite(inertias).all()
or not np.isfinite(friction).all()
or np.any(friction < 0.0)
):
raise RuntimeError(f"entity {entity.name} native property readback is invalid")
record: dict[str, Any] = {
"name": entity.name,
"body_mass": masses.tolist(),
"body_com": coms[:, :, :3].tolist(),
"body_inertia": inertias.tolist(),
"geom_friction": friction.tolist(),
}
if entity.kind == "articulation" and entity.joints:
for field, native in (
("dof_stiffness", asset.root_physx_view.get_dof_stiffnesses()),
("dof_damping", asset.root_physx_view.get_dof_dampings()),
("dof_armature", asset.root_physx_view.get_dof_armatures()),
("dof_friction", self._native_dof_friction(asset)),
):
table = _numpy(native)[mapping["envs"]][:, mapping["joints"]]
if not np.isfinite(table).all():
raise RuntimeError(
f"entity {entity.name} native {field} readback is invalid"
)
record[field] = table.tolist()
records.append(record)
return records
def _verify_reset_property_readback(
self,
ids: np.ndarray,
randomization: dict[str, np.ndarray],
before: list[dict[str, Any]],
records: list[dict[str, Any]],
) -> None:
selected = np.zeros(self.num_envs, dtype=bool)
selected[ids] = True
def check(
entity: Any,
field: str,
previous: dict[str, Any],
record: dict[str, Any],
expected: np.ndarray,
) -> None:
actual = np.asarray(record[field], dtype=np.float32).reshape(
self.num_envs, *expected.shape[1:]
)
if not np.allclose(actual[ids], expected, rtol=2e-5, atol=1e-6):
raise RuntimeError(
f"entity {entity.name} native {field} readback differs from reset: "
f"expected {expected.tolist()}, got {actual[ids].tolist()}"
)
untouched = np.asarray(previous[field], dtype=np.float32).reshape(actual.shape)
if not np.allclose(actual[~selected], untouched[~selected], rtol=2e-5, atol=1e-6):
raise RuntimeError(
f"entity {entity.name} native {field} write leaked outside selected rows"
)
geom_offset = 0
for index, (entity, previous, record) in enumerate(
zip(self.layout.entities, before, records)
):
count = len(entity.geoms)
if "geom_friction" in randomization:
# Records use the full public layout. Visual-only rows have no
# native shape, so they retain their previous values.
valid = self.maps[index]["geoms"] >= 0
full_previous = np.asarray(previous["geom_friction"], dtype=np.float32).reshape(
self.num_envs, count, 3
)
full_expected = full_previous.copy()
full_expected[ids[:, None], valid] = randomization["geom_friction"][
:, geom_offset : geom_offset + count, :
][:, valid, :]
check(
entity,
"geom_friction",
previous,
record,
full_expected[ids],
)
geom_offset += count
if "body_mass" in randomization:
check(
entity,
"body_mass",
previous,
record,
randomization["body_mass"][:, list(entity.body_ids)],
)
if "body_ipos" in randomization:
check(
entity,
"body_com",
previous,
record,
randomization["body_ipos"][:, list(entity.body_ids), :],
)
if "body_inertia" in randomization:
expected = np.asarray(
[
self._expected_inertia_blocks(index, int(env), row)
for env, row in zip(ids, randomization["body_inertia"])
]
)
check(entity, "body_inertia", previous, record, expected)
self._verify_dof_readback(entity, index, ids, randomization, previous, record, check)
def _verify_dof_readback(
self,
entity: Any,
index: int,
ids: np.ndarray,
randomization: dict[str, np.ndarray],
previous: dict[str, Any],
record: dict[str, Any],
check: Any,
) -> None:
if entity.kind != "articulation" or not entity.joints:
return
joint_columns = [joint.qvel_indices[0] for joint in entity.joints]
actuator_positions = [
position
for position, joint in enumerate(entity.joints)
if joint.name in entity.actuator_joint_names
]
kp = randomization.get("kp")
if kp is not None and actuator_positions:
expected = np.asarray(previous["dof_stiffness"], dtype=np.float32)[ids].copy()
expected[:, actuator_positions] = kp[:, list(entity.actuator_indices)]
check(entity, "dof_stiffness", previous, record, expected)
if "kd" in randomization or "dof_damping" in randomization:
drive = self._drive_damping[index]
natural = self._natural_damping[index]
assert drive is not None and natural is not None
check(entity, "dof_damping", previous, record, drive[ids] + natural[ids])
if "dof_armature" in randomization:
check(
entity,
"dof_armature",
previous,
record,
randomization["dof_armature"][:, joint_columns],
)
if "dof_frictionloss" in randomization:
check(
entity,
"dof_friction",
previous,
record,
randomization["dof_frictionloss"][:, joint_columns],
)
[docs]
def reset_entities(self, payload: dict[str, Any]) -> dict[str, Any]:
count = payload["count"]
if isinstance(count, bool) or not isinstance(count, int) or not 0 < count <= self.num_envs:
raise ValueError("invalid reset count")
ids = self.slots["reset_env_ids"][:count].astype(np.int64, copy=True)
if np.unique(ids).size != count or np.any(ids < 0) or np.any(ids >= self.num_envs):
raise ValueError("invalid reset environment IDs")
names = payload["entity_names"]
if not isinstance(names, list) or len(names) != len(set(names)):
raise ValueError("entity_names must be unique names")
if not set(names).issubset({entity.name for entity in self.layout.entities}):
raise ValueError("unknown reset entity")
qpos = self.slots["reset_qpos"][:count].copy()
qvel = self.slots["reset_qvel"][:count].copy()
roots = self.slots["reset_entity_root_state"][:count].copy()
pmask = self.slots["reset_qpos_mask"].copy()
vmask = self.slots["reset_qvel_mask"].copy()
rmask = self.slots["reset_root_mask"].copy()
if any(not np.isfinite(values).all() for values in (qpos, qvel, roots)):
raise ValueError("reset values must be finite")
if any(np.any((mask != 0) & (mask != 1)) for mask in (pmask, vmask, rmask)):
raise ValueError("reset masks must contain zero or one")
control_values = None
control_columns = tuple(
column
for entity in self.layout.entities
if entity.name in names
for column in entity.actuator_indices
)
if "control_values" in payload:
raw_control = np.asarray(payload["control_values"])
if (
raw_control.shape != (count, self.layout.nu)
or raw_control.dtype.kind not in "fiu"
or not np.isfinite(raw_control).all()
or np.any(np.abs(raw_control.astype(np.float64)) > np.finfo(np.float32).max)
):
raise ValueError("control_values must be finite (count, nu) float32 values")
control_values = raw_control.astype(np.float32, copy=True)
unselected = [
column for column in range(self.layout.nu) if column not in control_columns
]
if not np.array_equal(
control_values[:, unselected], self.slots["ctrl"][ids][:, unselected]
):
raise ValueError("control_values changes controls of an unselected entity")
for index, entity in enumerate(self.layout.entities):
selected = bool(
pmask[list(entity.qpos_indices)].any()
or vmask[list(entity.qvel_indices)].any()
or rmask[index].any()
)
if selected and entity.name not in names:
raise ValueError("reset mask writes undeclared entity")
if entity.root_mode == "fixed" and rmask[index].any():
raise ValueError("fixed root cannot be reset")
if entity.root_mode == "kinematic" and rmask[index, 1]:
raise ValueError("kinematic root velocity cannot be reset")
if rmask[index, 0] and not np.allclose(
np.linalg.norm(roots[:, index, 3:7], axis=1), 1, rtol=0, atol=1e-5
):
raise ValueError("root reset quaternion must be unit wxyz")
if entity.root_mode == "floating":
for cols, mask, channel in (
(entity.root_qpos_indices, pmask, 0),
(entity.root_qvel_indices, vmask, 1),
):
bits = mask[list(cols)]
if np.any(bits) != bool(rmask[index, channel]) or len(set(bits.tolist())) > 1:
raise ValueError("root generalized and entity masks disagree")
if rmask[index, 0] and not np.allclose(
qpos[:, entity.root_qpos_indices], roots[:, index, :7], rtol=1e-5, atol=1e-6
):
raise ValueError("root pose differs between generalized and entity values")
if rmask[index, 1]:
velocity = roots[:, index, 7:].copy()
velocity[:, 3:] = _rotate(roots[:, index, 3:7], velocity[:, 3:], inverse=True)
if not np.allclose(
qvel[:, entity.root_qvel_indices], velocity, rtol=1e-5, atol=1e-6
):
raise ValueError("generalized and entity root velocities differ")
randomization = self._validated_reset_randomization(payload, count)
self._commit(ids, qpos, qvel, roots, pmask, vmask, rmask)
try:
property_records = None
if randomization is not None:
before_properties = self._readback_reset_properties()
self._apply_reset_randomization(ids, randomization)
property_records = self._readback_reset_properties()
self._verify_reset_property_readback(
ids, randomization, before_properties, property_records
)
if control_values is not None:
for entity, asset, mapping in zip(self.layout.entities, self.assets, self.maps):
if entity.name in names and entity.actuator_indices:
native_ids = self.torch.as_tensor(
mapping["envs"][ids], dtype=self.torch.long, device=self.device
)
asset.set_joint_position_target(
self._tensor(control_values[:, entity.actuator_indices]),
joint_ids=mapping["controls"].tolist(),
env_ids=native_ids,
)
self.slots["ctrl"][np.ix_(ids, control_columns)] = control_values[
:, control_columns
]
self.refresh_state_slots()
except Exception:
self.faulted = True
raise
if randomization is not None:
assert property_records is not None
for entity, record in zip(self.layout.entities, property_records):
current = (
self.actual[entity.name]
if isinstance(self.actual, dict)
else next(item for item in self.actual if item["name"] == entity.name)
)
for field in (
"body_mass",
"body_com",
"body_inertia",
"geom_friction",
):
if field in record and field in current:
current[field] = record[field]
return {"timing": {}, "native_entity_records": property_records}
return {"timing": {}}
[docs]
def get_meta(self) -> dict[str, Any]:
if self._model_file_metadata is not None:
return self._model_file_metadata.copy()
effective: dict[str, Any] = {
"dt": float(self.sim.get_physics_dt()),
"share_friction_materials": getattr(self, "_share_friction_materials", False),
"env_spacing": getattr(self, "_env_spacing", 2.0),
"gravity": self.gravity.tolist(),
"collision_filter": {
"self_collision": {
entry["name"]: bool(entry["self_collision"]) for entry in self.entries
},
"environment_isolation": True,
"implicit_ground": False,
},
"entity_gravity_disabled": {
entry["name"]: bool(entry["gravity_disabled"]) for entry in self.entries
},
# Verified at INIT against the PhysX max-velocity readback; only
# entities with an explicit declaration appear.
"entity_joint_velocity_limits": {
entry["name"]: dict(entry["joint_velocity_limits"])
for entry in self.entries
if entry.get("joint_velocity_limits") is not None
},
}
engine_readback = ["dt"]
solver_fields = self.physx_solver.configured_fields()
if solver_fields:
engine_fields = [
field for field in solver_fields if field not in PHYSX_SOLVER_AUTHORED_FIELDS
]
solver_readback = (
read_engine_solver_values(
self.sim.stage,
include_contact_offset=self.physx_solver.contact_offset is not None,
include_rest_offset=self.physx_solver.rest_offset is not None,
include_max_depenetration_velocity=(
self.physx_solver.max_depenetration_velocity is not None
),
)
if engine_fields
else {}
)
for field in solver_fields:
if field in PHYSX_SOLVER_AUTHORED_FIELDS:
# Carb settings have no USD attribute; report the authored
# value instead of an engine readback.
effective[field] = getattr(self.physx_solver, field)
else:
effective[field] = solver_readback[field]
engine_readback.append(field)
return {
"scene_layout": self.layout.to_dict(),
"scene_entities_actual": self.actual,
"gravity": self.gravity.tolist(),
"use_gpu_pipeline": True,
"env_origins": self.origins.tolist(),
"collision_filtering_applied": self.num_envs > 1,
"render_mode": self.renderer.render_mode,
"render_width": self.renderer.render_width,
"render_height": self.renderer.render_height,
"graphics_enabled": self.renderer.render_mode != "none",
"share_friction_materials": getattr(self, "_share_friction_materials", False),
"raw_usd_cache": {
"enabled": self._raw_usd_cache_persistent,
"unique_sources": len(self._reported_raw_usd_source_digests),
"hits": sum(item["hit"] for item in self._raw_usd_cache_reports),
"conversions": sum(not item["hit"] for item in self._raw_usd_cache_reports),
"entries": tuple(self._raw_usd_cache_reports),
},
"role_usd_cache": {
"enabled": self._role_usd_cache is not None,
"hits": sum(item["hit"] for item in self._role_usd_cache_reports),
"bakes": sum(not item["hit"] for item in self._role_usd_cache_reports),
"entries": tuple(self._role_usd_cache_reports),
},
"configuration_report": {
"schema_version": 1,
"effective": effective,
"engine_readback": engine_readback,
},
# Additive diagnostics key: mapped-scene INIT phase timings and
# resource snapshots (None until init_sim completes).
"init_telemetry": self._init_telemetry,
}
[docs]
def shutdown(self) -> None:
cuda_ipc, self._cuda_ipc = self._cuda_ipc, None
self._cuda_maps = []
if cuda_ipc is not None:
try:
cuda_ipc.close()
except Exception:
pass
for handle in self._shm_handles:
handle.close()
self._shm_handles.clear()
self.renderer.shutdown()
self._temporary.cleanup()
[docs]
def bind_renderer(self, payload: dict[str, Any]) -> dict[str, Any]:
return cast(dict[str, Any], self.renderer.bind_renderer(payload))
[docs]
def render_frame(self) -> dict[str, Any]:
return cast(dict[str, Any], self.renderer.render_frame())
[docs]
def capture_frame(self) -> dict[str, Any]:
return cast(dict[str, Any], self.renderer.capture_frame())