"""IsaacGym specialization of the shared MJCF subprocess host adapter.
IsaacGym Preview 4 runs in an external Python 3.8 process. The shared owner
layer supplies pipe/shm lifecycle, MJCF metadata, selected reset, NumPy state
views, and native-render command plumbing; this module only selects the
IsaacGym runtime and worker entrypoint.
"""
from __future__ import annotations
import math
from collections.abc import Callable, Mapping
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any
import numpy as np
from unisim.backend.base import (
PhysicsStateLayout,
PlayCapabilities,
PublicStateWidths,
SelectedResetPublication,
TensorDataPlane,
TensorExecution,
TensorLifecycleCapabilities,
TensorProcessTopology,
validate_torch_device,
)
from unisim.backend.mjcf_layout import extract_mjcf_joint_layout
from unisim.backend.subprocess_ipc import protocol
from unisim.backend.subprocess_ipc.backend import (
_ROOT_QPOS_DIM,
_ROOT_QVEL_DIM,
MjcfSubprocessBackend,
SubprocessModelInfo,
SubprocessWorkerError,
_normalize_camera_kwargs,
)
from unisim.backend.subprocess_ipc.isaac_common import (
mapped_default_actuator_rows,
mapped_default_body_rows,
mapped_default_dof_rows,
mapped_default_full_geom_friction,
)
from unisim.dr._host import interval_payload_to_host, reset_payload_to_host_dict
from unisim.dr.interval import (
INTERVAL_TERM_BODY_FORCE,
INTERVAL_TERM_BODY_TORQUE,
IntervalTermOp,
)
from unisim.dr.types import (
RESET_TERM_BODY_INERTIA,
RESET_TERM_BODY_IPOS,
RESET_TERM_BODY_MASS,
RESET_TERM_DOF_ARMATURE,
RESET_TERM_DOF_FRICTIONLOSS,
RESET_TERM_GEOM_FRICTION,
RESET_TERM_KD,
RESET_TERM_KP,
DomainRandomizationCapabilities,
FixedVariantLayout,
IntervalRandomizationPlan,
TensorResetRandomizationPayload,
_validate_reset_term,
validate_tensor_reset_payload,
)
from unisim.entities import SceneResetRequest
from unisim.inspection import ConfigurationField, ConfigurationProvenance
from unisim.progress import progress_enabled
from unisim.utils.quaternion import torch_quat_mul, torch_quat_rotate
from .dependencies import build_worker_env, resolve_isaacgym_runtime
from .tensor import IsaacGymCudaIpcPlan
_WORKER_PATH = Path(__file__).resolve().parent / "worker.py"
_MAPPED_RESET_TERMS = frozenset(
{
RESET_TERM_KP,
RESET_TERM_KD,
RESET_TERM_BODY_MASS,
RESET_TERM_BODY_INERTIA,
RESET_TERM_BODY_IPOS,
RESET_TERM_DOF_ARMATURE,
RESET_TERM_DOF_FRICTIONLOSS,
RESET_TERM_GEOM_FRICTION,
}
)
_MAPPED_INTERVAL_TERMS = frozenset({INTERVAL_TERM_BODY_FORCE, INTERVAL_TERM_BODY_TORQUE})
_FIXED_VARIANT_LAYOUTS = frozenset(
{
FixedVariantLayout.SAME_LAYOUT,
FixedVariantLayout.UNIFORM_PUBLIC_LAYOUT,
}
)
[docs]
class IsaacGymWorkerError(SubprocessWorkerError):
"""Raised when the external IsaacGym worker fails."""
[docs]
@dataclass(frozen=True)
class IsaacGymModelInfo(SubprocessModelInfo):
"""Opaque metadata returned by the IsaacGym worker handshake."""
[docs]
class IsaacGymBackend(MjcfSubprocessBackend):
"""NumPy-facing ``SimBackend`` client for the IsaacGym worker."""
_BACKEND_TYPE = "isaacgym"
_BACKEND_LABEL = "isaacgym"
_WORKER_ERROR_CLS = IsaacGymWorkerError
_MODEL_INFO_CLS = IsaacGymModelInfo
[docs]
def __init__(
self,
scene: Any,
num_envs: int,
sim_dt: float,
*,
env_spacing: float = 4.0,
**kwargs: Any,
) -> None:
if isinstance(env_spacing, bool) or not isinstance(env_spacing, (int, float)):
raise ValueError(f"env_spacing must be positive and finite, got {env_spacing!r}")
if not math.isfinite(float(env_spacing)) or env_spacing <= 0:
raise ValueError(f"env_spacing must be positive and finite, got {env_spacing!r}")
self._env_spacing = float(env_spacing)
requested_device_id = kwargs.get("device_id", 0)
if isinstance(requested_device_id, bool) or not isinstance(requested_device_id, int):
raise TypeError(f"device_id must be an integer or None, got {requested_device_id!r}")
if requested_device_id < -1:
raise ValueError(f"device_id must be -1 or >= 0, got {requested_device_id!r}")
self._requested_device_id = requested_device_id
super().__init__(scene, num_envs, sim_dt, **kwargs)
self._staged_body_wrench = (
None
if self._entity_scene is None
else np.zeros((self._num_envs, self._entity_scene.layout.nbody, 6), dtype=np.float32)
)
self._body_wrench_pending = False
self._cuda_ipc_plan: IsaacGymCudaIpcPlan | None = None
# Physics-state playback support state. The worker protocol carries no
# simulation clock, so the snapshot time column is a per-env host
# accumulator advanced by nsteps * sim_dt on every public step path and
# zeroed on state resets (mirroring the MuJoCo/Motrix time semantics).
self._time_view = np.zeros((self._num_envs,), dtype=np.float64)
self._playback_unsupported_reason: str | None = None
self._playback_state_supported = (
True if self._entity_scene is not None else self._audit_model_file_playback_shape()
)
def _audit_model_file_playback_shape(self) -> bool:
"""Audit a raw-MJCF scene for physics-state playback at construction.
The model-file worker wire always synthesizes a 7/6 floating root plus one
column per single-DoF joint, so a playback snapshot is only meaningful
to a MuJoCo playback shell when the source MJCF declares exactly one
free joint as its first joint and every remaining joint is a named
hinge/slide in matching document order. Shapes that conflict with the
model-file wire itself (multiple free roots, non-single-DoF joints,
unnamed joints) fail closed here; shapes whose physics wire is
well-defined but whose playback shell cannot match the synthetic root
(no free joint, or a free joint that is not the first joint) disable
the capability instead, and the layout getter raises with the reason.
"""
model_file = str(Path(self._scene.model_file).expanduser())
entries = extract_mjcf_joint_layout(model_file)
free = [entry for entry in entries if entry.kind == "free"]
joints = [entry for entry in entries if entry.kind != "free"]
if len(free) > 1:
raise NotImplementedError(
"isaacgym physics-state playback requires at most one MJCF free joint; "
f"{model_file} declares {len(free)}"
)
for entry in joints:
if entry.kind not in ("hinge", "slide"):
raise NotImplementedError(
"isaacgym physics-state playback requires single-DoF MJCF joints; "
f"joint {entry.name or entry.body_name!r} on body {entry.body_name!r} "
f"is a {entry.kind} joint (the model-file worker wire carries one dof "
"per joint)"
)
if not entry.name:
raise NotImplementedError(
"isaacgym physics-state playback requires named MJCF joints; a "
f"{entry.kind} joint on body {entry.body_name!r} is unnamed"
)
if not free:
self._playback_unsupported_reason = (
"the raw MJCF scene declares no free joint, so a MuJoCo playback "
"shell compiled from it cannot match the synthetic 7/6 floating "
"root of the model-file worker wire"
)
return False
if entries[0] is not free[0]:
self._playback_unsupported_reason = (
"the raw MJCF free joint is not the first joint in document order, "
"so the MuJoCo playback shell's generalized-state order cannot "
"match the model-file worker wire"
)
return False
return True
def _validate_model_file_playback_joint_order(self) -> None:
"""Fail closed when the native dof order differs from the MJCF source order.
Raw-MJCF playback snapshots serve the worker's canonical
``qpos``/``qvel`` rows, whose dof columns follow the native asset dof
order; those columns are only meaningful to a MuJoCo playback shell
while the native order matches the source MJCF joint order. This
mirrors the Motrix non-portable validation: the public addresses are
the synthetic 7/6 root plus the MJCF single-DoF joints in document
order, checked against the worker-reported dof names at handshake.
"""
if not self._playback_state_supported:
return
assert self._model_info is not None
entries = extract_mjcf_joint_layout(str(Path(self._scene.model_file).expanduser()))
joints = [entry for entry in entries if entry.kind != "free"]
names = tuple(entry.name for entry in joints)
if names != self._model_info.dof_names:
raise self._worker_error(
"isaacgym native dof order differs from the MJCF source joint order:\n"
f" mjcf: {names}\n"
f" worker: {self._model_info.dof_names}\n"
"Physics-state playback snapshots would not align with a MuJoCo "
"playback shell compiled from the scene source."
)
for index, entry in enumerate(joints):
if (
entry.qpos_address != _ROOT_QPOS_DIM + index
or entry.qvel_address != _ROOT_QVEL_DIM + index
or entry.num_dof_pos != 1
or entry.num_dof_vel != 1
):
raise self._worker_error(
f"isaacgym MJCF joint {entry.name!r} generalized-state addresses "
f"({entry.qpos_address}, {entry.qvel_address}) differ from the "
"model-file worker wire layout"
)
def _worker_init_payload(self) -> dict[str, Any]:
# Progress frames interleave with the INIT reply; only opt in when the
# host will render them so raw-protocol clients never see them.
return {"env_spacing": self._env_spacing, "init_progress": progress_enabled()}
def _resolve_worker_entity_gravity(self, prepared: Any) -> None:
"""IsaacGym historically kept gravity enabled on every entity asset."""
for entry in prepared.payload["scene_entities"]:
if entry["gravity_disabled"] is None:
entry["gravity_disabled"] = False
def _validate_reported_entity_gravity_disabled(self, meta: dict[str, Any]) -> None:
"""Strictly compare reported per-entity gravity state with the INIT request."""
assert self._entity_scene is not None
envelope = meta.get("configuration_report")
effective = envelope.get("effective") if isinstance(envelope, dict) else None
reported = effective.get("entity_gravity_disabled") if isinstance(effective, dict) else None
expected = {
entry["name"]: bool(entry["gravity_disabled"])
for entry in self._entity_scene.payload["scene_entities"]
}
if not isinstance(reported, dict) or set(reported) != set(expected):
raise self._worker_error(
"isaacgym worker did not report per-entity gravity state for the "
f"mapped scene: worker={reported!r}, host={expected!r}"
)
mismatched = sorted(
name
for name, wanted in expected.items()
if not isinstance(reported[name], bool) or reported[name] != wanted
)
if mismatched:
raise self._worker_error(
"isaacgym worker per-entity gravity state does not match the host INIT "
f"request for entities: {', '.join(mismatched)}"
)
def _validate_reported_entity_self_collision(self, meta: dict[str, Any]) -> None:
"""Strictly compare reported per-entity self-collision with the INIT request."""
assert self._entity_scene is not None
envelope = meta.get("configuration_report")
effective = envelope.get("effective") if isinstance(envelope, dict) else None
collision_filter = (
effective.get("collision_filter") if isinstance(effective, dict) else None
)
reported = (
collision_filter.get("self_collision") if isinstance(collision_filter, dict) else None
)
expected = {
entry["name"]: bool(entry["self_collision"])
for entry in self._entity_scene.payload["scene_entities"]
}
if not isinstance(reported, dict) or set(reported) != set(expected):
raise self._worker_error(
"isaacgym worker did not report per-entity self-collision for the "
f"mapped scene: worker={reported!r}, host={expected!r}"
)
mismatched = sorted(
name
for name, wanted in expected.items()
if not isinstance(reported[name], bool) or reported[name] != wanted
)
if mismatched:
raise self._worker_error(
"isaacgym worker per-entity self-collision does not match the host "
f"INIT request for entities: {', '.join(mismatched)}"
)
def _bind_model_metadata(self, meta: dict[str, Any]) -> None:
super()._bind_model_metadata(meta)
if self._entity_scene is None:
self._validate_model_file_playback_joint_order()
def _bind_scene_metadata(self, meta: dict[str, Any]) -> None:
super()._bind_scene_metadata(meta)
reported_spacing = meta.get("env_spacing")
if (
isinstance(reported_spacing, bool)
or not isinstance(reported_spacing, (int, float))
or not math.isfinite(float(reported_spacing))
or not math.isclose(
float(reported_spacing), self._env_spacing, rel_tol=0.0, abs_tol=1e-9
)
):
raise self._worker_error(
f"worker env spacing differs from the requested {self._env_spacing}: "
f"{reported_spacing!r}"
)
origins = np.asarray(meta.get("env_origins"), dtype=float)
columns = max(1, math.ceil(math.sqrt(self._num_envs)))
indexes = np.arange(self._num_envs)
expected = (
np.column_stack(
(
indexes % columns,
indexes // columns,
np.zeros(self._num_envs, dtype=float),
)
)
* self._env_spacing
)
if (
origins.shape != (self._num_envs, 3)
or not np.isfinite(origins).all()
or not np.allclose(origins, expected, rtol=0.0, atol=1e-6)
):
raise self._worker_error(
f"worker env origins do not match a {self._env_spacing:g} m native grid: "
f"{origins.tolist()!r}"
)
if self._entity_scene is not None:
self._validate_reported_entity_gravity_disabled(meta)
self._validate_reported_entity_self_collision(meta)
def _capture_entity_report(self, meta: dict[str, Any]) -> None:
super()._capture_entity_report(meta)
report = self.get_import_report()
spacing = ConfigurationField(
"env_spacing",
self._env_spacing,
float(meta["env_spacing"]),
"exact",
(
ConfigurationProvenance("adapter_setting", "Native worker scene profile"),
ConfigurationProvenance("engine_readback", "Worker-reported native env origins"),
),
unit="m",
reason="Adjacent environment origins are audited against the requested grid.",
)
self._import_report = replace(report, fields=(*report.fields, spacing))
def _supports_fixed_variant_plans(self) -> bool:
return True
def _worker_entrypoint(self) -> Path:
return _WORKER_PATH
def _resolve_worker_runtime(self) -> Any:
return resolve_isaacgym_runtime()
def _build_worker_environment(self, runtime: Any) -> dict[str, str]:
return build_worker_env(runtime)
def _runtime_payload(self, runtime: Any) -> dict[str, str]:
return {"isaacgym_python": str(runtime.isaacgym_python)}
[docs]
def get_dr_capabilities(self) -> DomainRandomizationCapabilities:
"""Advertise mapped-scene reset DR and interval wrenches.
The model-file raw-MJCF path declares only the selected-reset PD-gain
bridge; mapped entity scenes declare the reset terms the worker writes
and reads back natively. Per-env ``gravity``/``dof_damping`` have no PhysX channel
(sim-wide parameter, drive-only joint damping) and the multi-root scene
has no single ``base`` target, so those terms stay unsupported. The
host-plane sets cover ``torch.cpu`` payloads; the device plane
(``torch.cuda``) carries the CUDA IPC PD-gain lane only.
"""
entity_plan = None if self._entity_scene is None else self._entity_scene.owner.variant_plan
has_plan = self._fixed_variant_plan is not None or entity_plan is not None
variant_kwargs: dict[str, Any] = (
{
"supports_fixed_variants": True,
"supported_fixed_variant_layouts": _FIXED_VARIANT_LAYOUTS,
"supports_per_env_playback": True,
}
if has_plan
else {}
)
device_reset_terms = frozenset({RESET_TERM_KP, RESET_TERM_KD})
if self._entity_scene is None:
return DomainRandomizationCapabilities(
supported_reset_terms=frozenset({RESET_TERM_KP, RESET_TERM_KD}),
supported_device_reset_terms=device_reset_terms,
**variant_kwargs,
)
return DomainRandomizationCapabilities(
supported_reset_terms=_MAPPED_RESET_TERMS,
supported_interval_terms=_MAPPED_INTERVAL_TERMS,
supported_device_reset_terms=device_reset_terms,
**variant_kwargs,
)
@staticmethod
def _coerce_mapped_reset_field(values: Any, name: str, shape: tuple[int, ...]) -> np.ndarray:
if not isinstance(values, np.ndarray) or not np.issubdtype(values.dtype, np.floating):
raise TypeError(f"isaacgym {name} must be a floating NumPy array")
array = np.asarray(values, dtype=np.float32)
if array.shape != shape:
raise ValueError(f"isaacgym {name} must have shape {shape}, got {array.shape}")
if not np.isfinite(array).all() or np.any(
np.abs(array.astype(np.float64)) > np.finfo(np.float32).max
):
raise ValueError(f"isaacgym {name} must contain finite float32 values")
return array.copy()
def _validated_mapped_reset_randomization(
self, randomization: TensorResetRandomizationPayload | None, rows: np.ndarray
) -> TensorResetRandomizationPayload | None:
if randomization is None:
return None
if not isinstance(randomization, TensorResetRandomizationPayload):
raise TypeError(
"isaacgym reset randomization must be a TensorResetRandomizationPayload"
)
if randomization.is_empty():
return None
if self._entity_scene is None:
raise NotImplementedError(
"isaacgym reset domain randomization requires a mapped entity scene"
)
# Host plane: the mapped subprocess worker is a NumPy boundary, so the
# torch-first payload must carry torch.cpu tensors.
validate_tensor_reset_payload(randomization, device="cpu")
unsupported = self.get_dr_capabilities().get_unsupported_reset_terms(
randomization.requested_terms(), plane="host"
)
if unsupported:
requested = ", ".join(sorted(unsupported))
raise NotImplementedError(
f"isaacgym does not support reset domain randomization terms: {requested}."
)
layout = self._entity_scene.layout
count = int(rows.size)
coerce = self._coerce_mapped_reset_field
host_values = reset_payload_to_host_dict(randomization)
def field(name: str, shape: tuple[int, ...]) -> np.ndarray | None:
values = host_values.get(name)
return None if values is None else coerce(values, name, shape)
kp = field(RESET_TERM_KP, (count, layout.nu))
kd = field(RESET_TERM_KD, (count, layout.nu))
for name, gains in ((RESET_TERM_KP, kp), (RESET_TERM_KD, kd)):
if gains is not None and np.any(gains < 0.0):
raise ValueError(f"isaacgym {name} values must be nonnegative")
body_mass = field(RESET_TERM_BODY_MASS, (count, layout.nbody))
if body_mass is not None and np.any(body_mass <= 0.0):
raise ValueError("isaacgym body_mass values must be positive")
body_ipos = field(RESET_TERM_BODY_IPOS, (count, layout.nbody, 3))
body_inertia = field(RESET_TERM_BODY_INERTIA, (count, layout.nbody, 3))
if body_inertia is not None and np.any(body_inertia <= 0.0):
raise ValueError("isaacgym body_inertia values must be positive")
root_columns = sorted(
{column for entity in layout.entities for column in entity.root_qvel_indices}
)
dof_armature = field(RESET_TERM_DOF_ARMATURE, (count, layout.nv))
dof_frictionloss = field(RESET_TERM_DOF_FRICTIONLOSS, (count, layout.nv))
for name, values in (
(RESET_TERM_DOF_ARMATURE, dof_armature),
(RESET_TERM_DOF_FRICTIONLOSS, dof_frictionloss),
):
if values is None:
continue
if np.any(values < 0.0):
raise ValueError(f"isaacgym {name} values must be nonnegative")
if root_columns and np.any(values[:, root_columns] != 0.0):
raise ValueError(
f"isaacgym {name} must be zero on floating-root columns; "
"PhysX root state is not a driven DOF"
)
geom_friction = field(RESET_TERM_GEOM_FRICTION, (count, layout.ngeom, 3))
if geom_friction is not None:
if np.any(geom_friction < 0.0):
raise ValueError("isaacgym geom_friction values must be nonnegative")
if np.any(geom_friction[..., 0] != geom_friction[..., 1]) or np.any(
geom_friction[..., 2] != 0.0
):
raise ValueError(
"isaacgym geom_friction requires static == dynamic and a zero third column"
)
import torch
return TensorResetRandomizationPayload(
**{
term: torch.from_numpy(values)
for term, values in (
(RESET_TERM_KP, kp),
(RESET_TERM_KD, kd),
(RESET_TERM_BODY_MASS, body_mass),
(RESET_TERM_BODY_IPOS, body_ipos),
(RESET_TERM_BODY_INERTIA, body_inertia),
(RESET_TERM_DOF_ARMATURE, dof_armature),
(RESET_TERM_DOF_FRICTIONLOSS, dof_frictionloss),
(RESET_TERM_GEOM_FRICTION, geom_friction),
)
if values is not None
}
)
def _consume_entity_reset_randomization(self, response: Any) -> None:
"""Replace current native property records from the reset barrier."""
if self._entity_scene is None:
raise self._worker_error("isaacgym property mutation requires a mapped scene")
layout = self._entity_scene.layout
raw_records = response.get("native_entity_records") if isinstance(response, dict) else None
if not isinstance(raw_records, list):
raise self._worker_error(
"isaacgym worker omitted native records after property mutation"
)
records = {record.get("name"): record for record in raw_records if isinstance(record, dict)}
names = {entity.name for entity in layout.entities}
if len(records) != len(raw_records) or set(records) != names:
raise self._worker_error(
"isaacgym worker property-mutation records do not match the frozen entity layout"
)
for entity in layout.entities:
current = self._native_entity_records.get(entity.name)
if current is None:
raise self._worker_error("worker omitted native entity properties: " + entity.name)
reported = records[entity.name]
shapes = {
"body_mass": (self._num_envs, len(entity.body_ids)),
"body_ipos": (self._num_envs, len(entity.body_ids), 3),
"body_inertia": (self._num_envs, len(entity.body_ids), 3, 3),
"dof_stiffness": (self._num_envs, len(entity.joints)),
"dof_damping": (self._num_envs, len(entity.joints)),
"dof_armature": (self._num_envs, len(entity.joints)),
"dof_friction": (self._num_envs, len(entity.joints)),
"geom_friction": (self._num_envs, len(entity.geoms), 3),
}
for field, shape in shapes.items():
try:
values = np.asarray(reported.get(field), dtype=np.float64)
except (TypeError, ValueError) as exc:
raise self._worker_error(
f"worker native {field} is malformed for entity {entity.name}: "
f"expected shape {shape}"
) from exc
if values.shape != shape or not np.isfinite(values).all():
raise self._worker_error(
f"worker native {field} is malformed for entity {entity.name}: "
f"got shape {values.shape}, expected {shape}"
)
current[field] = reported[field]
# ------------------------------------------------------------------ #
# Mapped reset term defaults
# ------------------------------------------------------------------ #
[docs]
def get_reset_term_default(self, term: str) -> np.ndarray:
"""Return per-environment variant-assigned mapped reset defaults."""
if self._entity_scene is None:
_validate_reset_term(term)
if not self.get_dr_capabilities().supports_reset_term(term):
raise NotImplementedError(f"IsaacGymBackend does not support reset term {term!r}")
kp, kd = self.get_actuator_gains()
value = kp if term == RESET_TERM_KP else kd
result = np.array(value, dtype=np.float32, copy=True)
result.setflags(write=False)
return result
_validate_reset_term(term)
if not self.get_dr_capabilities().supports_reset_term(term):
raise NotImplementedError(f"IsaacGymBackend does not support reset term {term!r}")
rows = np.arange(self._num_envs, dtype=np.intp)
if term == RESET_TERM_BODY_MASS:
value = self._mapped_default_body_rows("body_mass", rows)
elif term == RESET_TERM_BODY_IPOS:
value = self._mapped_default_body_rows("body_ipos", rows, width=3)
elif term == RESET_TERM_BODY_INERTIA:
value = self._mapped_default_body_rows("body_inertia", rows, width=3)
elif term == RESET_TERM_GEOM_FRICTION:
value = self._mapped_default_geom_friction(rows)
elif term == RESET_TERM_KP:
value = self._mapped_default_actuator_rows("dof_stiffness", rows)
elif term == RESET_TERM_KD:
value = self._mapped_default_actuator_rows("dof_damping", rows)
elif term == RESET_TERM_DOF_ARMATURE:
value = self._mapped_default_dof_rows("dof_armature", rows)
elif term == RESET_TERM_DOF_FRICTIONLOSS:
value = self._mapped_default_dof_rows("dof_friction", rows)
else: # pragma: no cover - guarded by supports_reset_term
raise NotImplementedError(f"IsaacGymBackend does not support reset term {term!r}")
result = np.array(value, dtype=np.float32, copy=True)
result.setflags(write=False)
return result
def _canonical_body_table(self, field: str, width: int | None = None) -> np.ndarray:
assert self._entity_scene is not None
expected = (
(self._entity_scene.layout.nbody, 3)
if width == 3
else (self._entity_scene.layout.nbody,)
)
try:
canonical = np.asarray(getattr(self._entity_scene.owner.model, field), dtype=np.float32)
except (TypeError, ValueError) as exc:
raise self._worker_error(
f"compiled canonical {field} is malformed: expected shape {expected}"
) from exc
if canonical.shape != expected or not np.isfinite(canonical).all():
raise self._worker_error(
f"compiled canonical {field} is malformed: got shape {canonical.shape}, "
f"expected {expected}"
)
return canonical.copy()
def _mapped_default_body_rows(
self, field: str, rows: np.ndarray, width: int | None = None
) -> np.ndarray:
"""Build per-row variant-assigned public body-property defaults."""
return mapped_default_body_rows(
self, self._canonical_body_table(field, width), field, rows, width
)
def _mapped_default_geom_friction(self, rows: np.ndarray) -> np.ndarray:
"""Build per-row variant-assigned public geom-friction defaults."""
return mapped_default_full_geom_friction(self, rows)
def _mapped_default_actuator_rows(self, field: str, rows: np.ndarray) -> np.ndarray:
"""Build per-row variant-assigned public actuator-gain defaults."""
return mapped_default_actuator_rows(self, field, rows)
def _mapped_default_dof_rows(self, field: str, rows: np.ndarray) -> np.ndarray:
"""Build per-row variant-assigned public DOF-property defaults."""
return mapped_default_dof_rows(self, field, rows)
# ------------------------------------------------------------------ #
# Interval body wrenches (mapped scenes)
# ------------------------------------------------------------------ #
[docs]
def apply_interval_randomization(self, plan: IntervalRandomizationPlan) -> None:
if plan.is_empty():
return
if self._entity_scene is not None:
assert self._staged_body_wrench is not None
self._staged_body_wrench.fill(0.0)
self._body_wrench_pending = False
super().apply_interval_randomization(plan)
def _interval_term_handlers(self) -> dict[str, Callable[[IntervalTermOp], None]]:
if self._entity_scene is None:
return super()._interval_term_handlers()
return {
INTERVAL_TERM_BODY_FORCE: self._apply_interval_body_force,
INTERVAL_TERM_BODY_TORQUE: self._apply_interval_body_torque,
}
def _apply_interval_body_force(self, op: IntervalTermOp) -> None:
if op.body_ids is None:
raise ValueError("interval body force requires body_ids")
self.apply_body_force(op.body_ids, interval_payload_to_host(op.payload))
def _apply_interval_body_torque(self, op: IntervalTermOp) -> None:
if op.body_ids is None:
raise ValueError("interval body torque requires body_ids")
zero_force = np.zeros((self._num_envs, len(op.body_ids), 3), dtype=np.float32)
self.apply_body_force(
op.body_ids, zero_force, torque=interval_payload_to_host(op.payload)
)
[docs]
def apply_body_force(
self,
body_ids: np.ndarray,
force: np.ndarray,
torque: np.ndarray | None = None,
) -> None:
"""Stage a world-frame wrench for the next mapped-scene step.
The force acts at each target body's center of mass and the optional
torque is about that center. Repeated submissions accumulate until the
next successful step consumes them.
"""
if self._entity_scene is None:
super().apply_body_force(body_ids, force, torque)
return
self._require_state("interval body force perturbation")
assert self._staged_body_wrench is not None
body_ids_np = np.asarray(body_ids)
if (
body_ids_np.ndim != 1
or body_ids_np.dtype.kind not in "iu"
or np.issubdtype(body_ids_np.dtype, np.bool_)
):
raise ValueError("body_ids must be a one-dimensional integer array")
if np.any(body_ids_np < 0) or np.any(body_ids_np >= self._staged_body_wrench.shape[1]):
raise ValueError(
f"body ids must be in [0, {self._staged_body_wrench.shape[1]}), "
f"got range [{body_ids_np.min(initial=0)}, {body_ids_np.max(initial=0)}]"
)
body_ids_np = body_ids_np.astype(np.intp, copy=False)
expected_shape = (self._num_envs, body_ids_np.size, 3)
force_np = np.asarray(force, dtype=np.float32)
if force_np.shape != expected_shape:
raise ValueError(f"body force must have shape {expected_shape}, got {force_np.shape}")
if not np.isfinite(force_np).all():
raise ValueError("body force contains NaN or Inf")
torque_np = None
if torque is not None:
torque_np = np.asarray(torque, dtype=np.float32)
if torque_np.shape != expected_shape:
raise ValueError(
f"body torque must have shape {expected_shape}, got {torque_np.shape}"
)
if not np.isfinite(torque_np).all():
raise ValueError("body torque contains NaN or Inf")
for body_offset, body_id in enumerate(body_ids_np):
self._staged_body_wrench[:, int(body_id), 0:3] += force_np[:, body_offset, :]
if torque_np is not None:
self._staged_body_wrench[:, int(body_id), 3:6] += torque_np[:, body_offset, :]
self._body_wrench_pending = bool(
np.any(self._staged_body_wrench) or self._body_wrench_pending
)
def _step_payload(self, nsteps: int) -> dict[str, Any]:
payload = super()._step_payload(nsteps)
if self._entity_scene is not None and self._body_wrench_pending:
assert self._staged_body_wrench is not None
# NumPy pickle internals are not stable across the host and IsaacGym
# worker interpreter versions; raw C-order bytes are.
payload["body_wrench"] = self._staged_body_wrench.tobytes(order="C")
return payload
def _after_step(self, payload: dict[str, Any]) -> None:
if "body_wrench" in payload:
assert self._staged_body_wrench is not None
self._staged_body_wrench.fill(0.0)
self._body_wrench_pending = False
[docs]
def get_playback_model(self, env_index: int | None = None) -> Any:
"""Return the playback model source used by detached MuJoCo replay.
Whole-MJCF (raw) scenes replay through their construction source and
mapped entity scenes through the composed MJCF; fixed-variant scenes
follow the Drake contract — an explicit ``env_index`` is required and
the row's assigned variant source is returned.
"""
plan = self._fixed_variant_plan
if plan is None:
if self._entity_scene is not None:
return super().get_playback_model(env_index)
if env_index is not None:
if isinstance(env_index, bool) or not isinstance(env_index, int):
raise TypeError("env_index must be an integer or None")
if env_index < 0 or env_index >= self._num_envs:
raise IndexError(f"env_index must be in [0, {self._num_envs - 1}]")
return self._get_scene_metadata().model_file
if env_index is None:
raise ValueError("fixed-variant playback requires an explicit env_index")
if isinstance(env_index, bool) or not isinstance(env_index, int):
raise TypeError("env_index must be an integer or None")
if env_index < 0 or env_index >= self._num_envs:
raise IndexError(f"env_index must be in [0, {self._num_envs - 1}]")
variant_index = int(plan.assignment[env_index])
return plan.variants[variant_index].model_file
# ------------------------------------------------------------------ #
# Physics-state playback contract #
# ------------------------------------------------------------------ #
[docs]
def get_play_capabilities(self) -> PlayCapabilities:
# The native interactive renderer and video capture path is unchanged;
# physics-state playback adds the detached MuJoCo replay contract when
# the construction-time audit accepted the scene shape. The getattr
# default keeps the report fail-closed for instances that bypassed
# __init__ (test doubles built via __new__).
playback_supported = getattr(self, "_playback_state_supported", False)
return PlayCapabilities(
supports_native_interactive_renderer=True,
supports_physics_state_playback=playback_supported,
supports_native_video_capture=True,
supports_mocap_playback=playback_supported and self._nmocap() > 0,
)
def _nmocap(self) -> int:
"""Kinematic entity roots, compiled as mocap bodies in the playback MJCF."""
if self._entity_scene is None:
return 0
return sum(
1 for entity in self._entity_scene.layout.entities if entity.root_mode == "kinematic"
)
[docs]
def get_physics_state_layout(self) -> PhysicsStateLayout:
"""Return the snapshot layout; unsupported scene shapes fail closed."""
if not self._playback_state_supported:
assert self._playback_unsupported_reason is not None
raise NotImplementedError(
"isaacgym does not support physics-state playback for this scene: "
f"{self._playback_unsupported_reason}"
)
if self._entity_scene is None:
num_dof = self._num_dof()
return PhysicsStateLayout(nq=_ROOT_QPOS_DIM + num_dof, nv=_ROOT_QVEL_DIM + num_dof)
layout = self._entity_scene.layout
return PhysicsStateLayout(nq=int(layout.nq), nv=int(layout.nv), nmocap=self._nmocap())
[docs]
def get_physics_state(self) -> np.ndarray:
"""Assemble contract ``[time, qpos, qvel, (mocap)?]`` rows in MuJoCo layout.
The worker owns the physics state; one ``GET_PHYSICS_STATE`` request
returns the canonical ``[qpos, qvel, (mocap_pos, mocap_quat)?]`` block
for every environment in a single reply frame (no out-of-band channel
and no per-env round trips). The leading time column is the host-side
per-env accumulator because the worker protocol carries no simulation
clock.
"""
layout = self.get_physics_state_layout()
self._require_state("physics snapshot")
reply = self._request(protocol.CMD_GET_PHYSICS_STATE, None, expect=protocol.CMD_META)
block = self._decode_physics_state_block(reply, layout)
out = np.empty((self._num_envs, layout.state_width), dtype=np.float32)
out[:, 0] = self._time_view
out[:, 1:] = block
return out
def _decode_physics_state_block(self, reply: Any, layout: PhysicsStateLayout) -> np.ndarray:
"""Strictly decode the worker's ``[qpos, qvel, (mocap)?]`` reply block."""
expected = (self._num_envs, layout.state_width - 1)
if not isinstance(reply, dict):
raise self._worker_error("isaacgym worker physics-state reply is malformed")
shape = reply.get("shape")
raw = reply.get("state")
if (
not isinstance(shape, (list, tuple))
or len(shape) != 2
or any(isinstance(dim, bool) or not isinstance(dim, int) for dim in shape)
or tuple(shape) != expected
):
raise self._worker_error(
f"isaacgym worker physics-state block shape must be {list(expected)}, got {shape!r}"
)
nbytes = expected[0] * expected[1] * np.dtype(np.float32).itemsize
if not isinstance(raw, bytes) or len(raw) != nbytes:
raise self._worker_error(
"isaacgym worker physics-state block must be C-order float32 bytes "
f"of length {nbytes}"
)
block = np.frombuffer(raw, dtype=np.float32).reshape(expected)
if not np.isfinite(block).all():
raise self._worker_error("isaacgym worker physics-state block is not finite")
return block
[docs]
def set_physics_state(self, state: np.ndarray) -> None:
"""Restore a ``get_physics_state`` snapshot through the set_state path."""
layout = self.get_physics_state_layout()
state_array = np.asarray(state, dtype=np.float32)
expected = (self._num_envs, layout.state_width)
if state_array.shape != expected:
raise ValueError(
"isaacgym physics snapshot must use the "
"[time, qpos, qvel, (mocap_pos, mocap_quat)?] layout with shape "
f"{expected}, got {state_array.shape}"
)
parts = layout.split_state(state_array)
# The split columns are strided views; the reset slots expect
# contiguous rows. set_state also refreshes the worker and host
# caches, keeping state and sensor getters consistent.
self._host_set_state(
np.arange(self._num_envs, dtype=np.intp),
np.ascontiguousarray(parts.qpos),
np.ascontiguousarray(parts.qvel),
)
# set_state restarts the playback clock for its rows; the snapshot's
# time column is authoritative for the restore.
self._time_view[...] = parts.time
[docs]
def get_playback_mocap_state(self, env_index: int = 0) -> tuple[np.ndarray, np.ndarray]:
"""Return copied ``(mocap_pos, mocap_quat)`` rows for detached playback."""
layout = self.get_physics_state_layout()
if layout.nmocap == 0:
raise NotImplementedError("isaacgym scene has no kinematic (mocap) entities to replay")
if isinstance(env_index, bool) or not isinstance(env_index, int):
raise TypeError("env_index must be an integer")
if env_index < 0 or env_index >= self._num_envs:
raise IndexError(f"env_index must be in [0, {self._num_envs - 1}]")
parts = layout.split_state(self.get_physics_state())
assert parts.mocap_pos is not None and parts.mocap_quat is not None
return parts.mocap_pos[env_index].copy(), parts.mocap_quat[env_index].copy()
# ------------------------------------------------------------------ #
# Playback clock maintenance #
# ------------------------------------------------------------------ #
def _host_step(self, ctrl: np.ndarray, nsteps: int = 1) -> dict[str, dict[str, float]]:
if getattr(self, "_cuda_ipc_plan", None) is not None:
raise RuntimeError("NumPy step is unavailable while a CUDA IPC tensor plan is open")
result = super()._host_step(ctrl, nsteps)
self._time_view += int(nsteps) * self._sim_dt
return result
def _host_reset(self, env_ids: np.ndarray | None = None) -> None:
if getattr(self, "_cuda_ipc_plan", None) is not None:
raise RuntimeError("NumPy reset is unavailable while a CUDA IPC tensor plan is open")
super()._host_reset(env_ids)
ids = (
np.arange(self._num_envs, dtype=np.intp)
if env_ids is None
else np.asarray(env_ids, dtype=np.intp)
)
self._time_view[ids] = 0.0
def _host_set_state(
self,
env_indices: np.ndarray,
qpos: np.ndarray,
qvel: np.ndarray,
randomization: TensorResetRandomizationPayload | None = None,
) -> dict[str, dict[str, float]]:
if getattr(self, "_cuda_ipc_plan", None) is not None:
raise RuntimeError(
"host-plane state writes are unavailable while a CUDA IPC plan is open"
)
result = super()._host_set_state(env_indices, qpos, qvel, randomization)
rows = np.asarray(env_indices, dtype=np.intp)
if rows.ndim == 1 and rows.size:
self._time_view[rows] = 0.0
return result
[docs]
def reset_entities(self, request: SceneResetRequest) -> None:
if getattr(self, "_cuda_ipc_plan", None) is not None:
raise RuntimeError(
"host-plane entity reset is unavailable while a CUDA IPC plan is open"
)
super().reset_entities(request)
self._time_view[np.asarray(request.env_ids, dtype=np.intp)] = 0.0
# ------------------------------------------------------------------ #
# Experimental CUDA IPC tensor lifecycle
# ------------------------------------------------------------------ #
[docs]
def get_tensor_execution(self) -> TensorExecution:
info = self._model_info
if info is None:
return (
TensorExecution.DEVICE_RESIDENT
if self._requested_device_id >= 0
else TensorExecution.UNSUPPORTED
)
if not bool(getattr(info, "use_gpu_pipeline", False)):
return TensorExecution.UNSUPPORTED
return TensorExecution.DEVICE_RESIDENT
[docs]
def get_tensor_capabilities(self) -> TensorLifecycleCapabilities:
if self.get_tensor_execution() is TensorExecution.UNSUPPORTED:
return TensorLifecycleCapabilities(execution=TensorExecution.UNSUPPORTED)
return TensorLifecycleCapabilities(
execution=TensorExecution.DEVICE_RESIDENT,
state_views=True,
state_fields=frozenset({"qpos", "qvel"}),
sensor_views=True,
stepping=True,
selected_reset=True,
fixed_variants=False,
host_pre_step_control=False,
packed_host_bridge=False,
process_topology=TensorProcessTopology.EXTERNAL_WORKER,
data_plane=TensorDataPlane.CUDA_IPC,
stream_event_ownership=(
"collector records control/reset-ready and worker records state-ready; "
"caller owns Torch stream"
),
reset_randomization=True,
device_reset_randomization=True,
packed_reset=True,
packed_reset_fields=frozenset({RESET_TERM_KP, RESET_TERM_KD}),
mask_selected_reset=True,
torch_devices=("cuda",),
selected_reset_publication=SelectedResetPublication.AUTHORITATIVE_VIEWS,
requires_post_construction_publication_barrier=True,
)
[docs]
def compile_cuda_ipc_io(self, device: Any | None = None) -> IsaacGymCudaIpcPlan:
if getattr(self, "_cuda_ipc_plan", None) is not None:
raise RuntimeError("an IsaacGym CUDA IPC tensor plan is already open")
self._cuda_ipc_plan = IsaacGymCudaIpcPlan(self, device)
return self._cuda_ipc_plan
[docs]
def get_state_views(
self, fields: tuple[str, ...] | str | None = None, device: Any | None = None
) -> Mapping[str, Any]:
plan = getattr(self, "_cuda_ipc_plan", None)
if plan is None:
plan = self.compile_cuda_ipc_io(device)
requested = (
("qpos", "qvel")
if fields is None
else ((fields,) if isinstance(fields, str) else tuple(fields))
)
unknown = set(requested) - {"qpos", "qvel"}
if unknown:
raise KeyError(f"unknown isaacgym tensor state field(s): {sorted(unknown)}")
if device is not None:
validate_torch_device(
(str(plan.device),),
device,
current_device=plan.device_index,
label="IsaacGym CUDA IPC state views",
)
views = plan.get_state_views(requested)
return {name: views[name] for name in requested}
[docs]
def get_public_state_widths(self) -> PublicStateWidths:
"""Return the synthetic raw-MJCF public tensor state widths."""
layout = self.get_physics_state_layout()
return PublicStateWidths(nq=int(layout.nq), nv=int(layout.nv))
def _cuda_derived_sensor_specs(self) -> dict[str, tuple[Any, int]]:
"""Return frame sensors projectable from public CUDA body state."""
return {
name: mapped
for name, mapped in getattr(self, "_sensor_map", {}).items()
if mapped[0].kind in ("framepos", "framequat", "framezaxis")
}
def _cuda_derived_sensor_view(self, name: str, plan: Any) -> Any:
candidates = [
(qualified, mapped)
for qualified, mapped in self._cuda_derived_sensor_specs().items()
if qualified == name or qualified.endswith(f"/{name}")
]
if len(candidates) != 1:
return None
_, (spec, _body_id) = candidates[0]
import torch
body_name = spec.body_name
body_pos = plan.get_sensor_view(f"track_pos_w_{body_name}")
body_quat = plan.get_sensor_view(f"track_quat_w_{body_name}")
if spec.kind == "framepos":
local_pos = torch.as_tensor(
spec.local_pos, dtype=torch.float32, device=body_pos.device
).expand(body_pos.shape[0], 3)
return body_pos + torch_quat_rotate(torch, body_quat, local_pos)
local_quat = torch.as_tensor(
spec.local_quat, dtype=torch.float32, device=body_quat.device
).expand(body_quat.shape[0], 4)
world_quat = torch_quat_mul(torch, body_quat, local_quat)
if spec.kind == "framequat":
return world_quat
world_z = torch.zeros_like(body_pos)
world_z[:, 2] = 1.0
return torch_quat_rotate(torch, world_quat, world_z)
def _bind_sensor_data_reader(self, names: tuple[str, ...]) -> Callable[[], np.ndarray]:
"""Bind the cold NumPy shape contract over CUDA sensor views."""
plan = getattr(self, "_cuda_ipc_plan", None)
if plan is None:
plan = self.compile_cuda_ipc_io(None)
def read() -> np.ndarray:
return np.concatenate(
[plan.get_sensor_view(name).detach().cpu().numpy() for name in names],
axis=1,
)
return read
[docs]
def get_sensor_view(self, name: str, device: Any | None = None) -> Any:
derived_candidates = [
qualified
for qualified in self._cuda_derived_sensor_specs()
if qualified == name or qualified.endswith(f"/{name}")
]
if (
name not in ("pelvis_local_linvel", "torso_gyro")
and len(derived_candidates) != 1
and not any(
name.startswith(prefix)
for prefix in (
"track_pos_w_",
"track_quat_w_",
"track_linvel_w_",
"track_angvel_w_",
)
)
):
raise KeyError(f"unknown IsaacGym CUDA IPC tensor sensor {name!r}")
plan = getattr(self, "_cuda_ipc_plan", None)
if plan is None:
plan = self.compile_cuda_ipc_io(None)
if device is not None:
validate_torch_device(
(str(plan.device),),
device,
current_device=plan.device_index,
label="IsaacGym CUDA IPC sensor views",
)
derived = self._cuda_derived_sensor_view(name, plan)
if derived is not None:
return derived
return plan.get_sensor_view(name)
[docs]
def step_tensor(self, ctrl: Any, nsteps: int = 1) -> dict[str, dict[str, float]]:
plan = getattr(self, "_cuda_ipc_plan", None)
if plan is None:
plan = self.compile_cuda_ipc_io(None)
return plan.step_tensor(ctrl, nsteps)
def _cuda_tensor_reset_wire_payload(
self,
randomization: Any,
*,
row_count: int,
device: Any,
) -> dict[str, Any] | None:
"""Validate PD gains and cross the worker reset-property boundary."""
if randomization is None or randomization.is_empty():
return None
if not isinstance(randomization, TensorResetRandomizationPayload):
raise TypeError(
"isaacgym CUDA IPC reset randomization requires a TensorResetRandomizationPayload"
)
supported = {RESET_TERM_KP, RESET_TERM_KD}
unsupported = set(randomization.requested_terms()) - supported
if unsupported:
raise NotImplementedError(
"isaacgym CUDA IPC reset randomization does not support terms: "
+ ", ".join(sorted(unsupported))
)
import torch
wire: dict[str, Any] = {}
for term in supported:
value = getattr(randomization, term)
if value is None:
continue
expected = (row_count, self.num_actuators)
if not isinstance(value, torch.Tensor):
raise TypeError(f"isaacgym CUDA IPC reset {term} must be a torch.Tensor")
if value.dtype != torch.float32 or not bool(value.is_contiguous()):
raise TypeError(f"isaacgym CUDA IPC reset {term} must be contiguous float32")
if tuple(value.shape) != expected:
raise ValueError(
f"isaacgym CUDA IPC reset {term} must have shape {expected}, "
f"got {tuple(value.shape)}"
)
if value.device != torch.device(device):
raise ValueError(f"isaacgym CUDA IPC reset {term} device does not match arena")
if not bool(torch.isfinite(value).all().item()):
raise ValueError(f"isaacgym CUDA IPC reset {term} values must be finite")
if bool((value < 0.0).any().item()):
raise ValueError(f"isaacgym CUDA IPC reset {term} values must be nonnegative")
wire[term] = value
return wire or None
[docs]
def set_state_tensor(
self,
env_indices: Any,
qpos: Any,
qvel: Any,
randomization: Any | None = None,
) -> dict[str, dict[str, float]]:
if randomization is not None and not isinstance(
randomization, TensorResetRandomizationPayload
):
raise TypeError(
"IsaacGym CUDA IPC reset randomization requires a TensorResetRandomizationPayload"
)
plan = getattr(self, "_cuda_ipc_plan", None)
if plan is None:
plan = self.compile_cuda_ipc_io(None)
wire_randomization = (
None
if randomization is None
else self._cuda_tensor_reset_wire_payload(
randomization,
row_count=int(env_indices.shape[0]),
device=plan.device,
)
)
return plan.set_state_tensor(env_indices, qpos, qvel, wire_randomization)
[docs]
def set_state_tensor_mask(
self,
reset_mask: Any,
qpos: Any,
qvel: Any,
randomization: Any | None = None,
*,
full_refresh: bool = False,
) -> dict[str, dict[str, float]] | None:
"""Commit full-width state through a backend-owned mask projection."""
import torch
if full_refresh:
raise NotImplementedError("isaacgym CUDA IPC mask reset is a selected-row boundary")
plan = getattr(self, "_cuda_ipc_plan", None)
if plan is None:
plan = self.compile_cuda_ipc_io(None)
if not isinstance(reset_mask, torch.Tensor):
raise TypeError("isaacgym CUDA IPC reset mask must be a torch.Tensor")
if (
reset_mask.ndim != 1
or reset_mask.dtype != torch.bool
or not bool(reset_mask.is_contiguous())
):
raise TypeError("isaacgym CUDA IPC reset mask must be a contiguous 1-D bool tensor")
if tuple(reset_mask.shape) != (self._num_envs,):
raise ValueError(f"isaacgym CUDA IPC reset mask must have shape {(self._num_envs,)}")
if not isinstance(qpos, torch.Tensor) or not isinstance(qvel, torch.Tensor):
raise TypeError("isaacgym CUDA IPC reset states must be torch.Tensors")
expected_qpos = (self._num_envs, plan.arena.nq)
expected_qvel = (self._num_envs, plan.arena.nv)
if qpos.dtype != torch.float32 or qvel.dtype != torch.float32:
raise TypeError("isaacgym CUDA IPC reset states must be float32")
if tuple(qpos.shape) != expected_qpos:
raise ValueError(f"isaacgym CUDA IPC reset qpos must have shape {expected_qpos}")
if tuple(qvel.shape) != expected_qvel:
raise ValueError(f"isaacgym CUDA IPC reset qvel must have shape {expected_qvel}")
if qpos.device != reset_mask.device or qvel.device != reset_mask.device:
raise ValueError("isaacgym CUDA IPC mask reset tensors must share one device")
self._cuda_tensor_reset_wire_payload(
randomization,
row_count=self._num_envs,
device=plan.device,
)
if not bool(reset_mask.any().item()):
return {"timing": {"set_state_tensor_mask_ms": 0.0}}
rows = reset_mask.nonzero(as_tuple=False).flatten()
selected_payload: TensorResetRandomizationPayload | None = None
if randomization is not None and not randomization.is_empty():
selected_payload = TensorResetRandomizationPayload(
**{
term: getattr(randomization, term).index_select(0, rows).contiguous()
for term in (RESET_TERM_KP, RESET_TERM_KD)
if getattr(randomization, term) is not None
}
)
return self.set_state_tensor(
rows,
qpos.index_select(0, rows).contiguous(),
qvel.index_select(0, rows).contiguous(),
selected_payload,
)
[docs]
def bind_packed_reset(self, *, device: Any | None = None) -> Any:
"""Begin a packed transaction over the external CUDA IPC reset boundary."""
from unisim.backend.packed_reset import PackedResetTransaction
return PackedResetTransaction(self, device=device)
[docs]
def close(self) -> None:
plan = getattr(self, "_cuda_ipc_plan", None)
if plan is not None and not plan.closed:
plan.close()
self._cuda_ipc_plan = None
super().close()
__all__ = [
"IsaacGymBackend",
"IsaacGymModelInfo",
"IsaacGymWorkerError",
"_normalize_camera_kwargs",
]