"""Cold-path scene materialization for the independent ``genesis`` backend.
Genesis 1.3.3 drops three MJCF features UniLab relies on, as measured in
the Genesis 1.3.3 feasibility audit (#1372 §3): ``<keyframe>``,
the global ``<option>`` block, and the whole ``<sensor>`` block. This module
compensates on the cold path exactly as the report prescribes: keyframe and
sensor metadata are scanned once with the ``mujoco`` package (the isaacgym
parent-process metadata-scan pattern) and cached — hot paths never parse XML;
global options arrive as explicit owner fields; and ``gs.init`` host-side
effects (torch default device/dtype and RNG) are contained by
:func:`preserve_torch_globals` plus the one-session-per-process guard.
"""
from __future__ import annotations
import os
import tempfile
from contextlib import contextmanager
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any, Iterator
import numpy as np
from unisim.backend.scene_cleanup import TemporarySceneCleanup
from unisim.scene import SceneCfg
# Supported spellings for the explicit owner-YAML global options. Genesis
# defaults apply when the owner leaves them unset (REPORT #1372 §5.2).
_INTEGRATOR_ENUMS = {
"euler": "Euler",
"implicitfast": "implicitfast",
"approximate_implicitfast": "approximate_implicitfast",
}
_CONSTRAINT_SOLVER_ENUMS = {
"newton": "Newton",
"cg": "CG",
}
_FRICTION_CONE_ENUMS = {
"pyramidal": "pyramidal",
"elliptic": "elliptic",
}
# Contact ``data="found"`` equivalent: per-link net contact force magnitude
# threshold in newtons (REPORT #1372 §3.4: standing G1 foot reads ~138 N).
CONTACT_FOUND_FORCE_THRESHOLD_N = 1.0
[docs]
@dataclass(frozen=True)
class GenesisSensorPlan:
"""One MJCF sensor mapped onto a Genesis-readable equivalent.
``body_name`` is the owning link for site sensors and the robot-side geom
body for contact sensors. ``site_pos``/``site_quat`` (wxyz) are the local
site frame in the body frame; both are ``None`` for contact and portable
body frame fragment sensors. Contact sensors additionally retain both final
geom names and whether the common fragment form is exact pair netforce, so
portable validation can bind them without source-model array indices.
"""
name: str
kind: str
dim: int
body_name: str
object_name: str
reference_type: int
reference_id: int
site_pos: tuple[float, ...] | None
site_quat: tuple[float, ...] | None
contact_geom1_name: str | None = None
contact_geom2_name: str | None = None
contact_netforce: bool = False
object_kind: str = "site"
[docs]
@dataclass(frozen=True)
class GenesisPortableEntitySource:
"""Independent normalized MJCF sources for one public entity."""
name: str
model_files: tuple[str, ...]
metadata: tuple[GenesisModelMetadata, ...]
[docs]
@dataclass(frozen=True)
class GenesisPortableSources:
"""Cold-path Genesis inputs derived from the common portable compiler."""
entities: tuple[GenesisPortableEntitySource, ...]
cleanup_handle: Any
[docs]
def genesis_balanced_variant_mapping(num_variants: int, num_envs: int) -> np.ndarray:
"""Return Genesis 1.3.3's documented heterogeneous env/variant mapping.
Genesis dispatches variants in contiguous balanced blocks (the first
remainder-sized block gets one extra environment). This local owner-side
implementation deliberately does not call Genesis' private solver helper;
construction accepts an assignment only when it is exactly this mapping.
"""
if num_variants < 1:
raise ValueError("genesis heterogeneous variants require at least one source")
if num_envs < 1:
raise ValueError("genesis heterogeneous variants require at least one environment")
if num_envs >= num_variants:
base, extra = divmod(num_envs, num_variants)
return np.repeat(
np.arange(num_variants, dtype=np.int32),
np.r_[np.full(extra, base + 1), np.full(num_variants - extra, base)],
)
return np.arange(num_envs, dtype=np.int32)
[docs]
def validate_genesis_variant_assignment(assignment: np.ndarray, num_variants: int) -> np.ndarray:
"""Require the exact mapping Genesis will materialize natively."""
values = np.asarray(assignment, dtype=np.int32)
if values.ndim != 1:
raise ValueError("genesis variant assignment must be one-dimensional")
expected = genesis_balanced_variant_mapping(num_variants, values.size)
if not np.array_equal(values, expected):
raise ValueError(
"genesis heterogeneous variant assignment must exactly equal the native "
f"balanced mapping {expected.tolist()}; got {values.tolist()}. Reorder the "
"assignment until Genesis exposes a public owner-controlled mapping."
)
return values.copy()
[docs]
def prepare_genesis_portable_sources(
mujoco: Any, scene: SceneCfg, layout: Any
) -> GenesisPortableSources:
"""Normalize common portable sources into independent Genesis MJCF files.
The common compiler remains the sole layout, identity, and source
relationship authority. This owner adapter only serializes each already
normalized ``MjSpec`` and scans the resulting source for native binding.
"""
from unisim.mjcf_compiler import load_entity_source
if any(
entity.root_mode == "kinematic" and entity.mirror_of is None
for entity in scene.entity_assets
):
raise NotImplementedError("genesis portable kinematic entities are not yet supported")
entity_specs = {entity.name: entity for entity in scene.entity_assets}
binding = scene.entity_variant
directory = tempfile.TemporaryDirectory(prefix="unisim-genesis-entities-")
entities: list[GenesisPortableEntitySource] = []
try:
for entity_spec in scene.entity_assets:
owner = layout.get_entity(entity_spec.name)
source_spec = (
entity_specs[entity_spec.mirror_of]
if entity_spec.mirror_of is not None
else entity_spec
)
paths: list[str]
if binding is not None and source_spec.name == binding.target_entity:
if owner.kind != "rigid" or len(owner.body_names) != 1:
raise NotImplementedError(
"genesis heterogeneous variants are native only for single-link "
f"rigid entities; entity {owner.name!r} is {owner.kind} with "
f"{len(owner.body_names)} bodies"
)
paths = [str(item.model_file) for item in binding.plan.variants]
else:
assert source_spec.source is not None
paths = [str(source_spec.source.model_file)]
files: list[str] = []
metadata: list[GenesisModelMetadata] = []
for variant, source_path in enumerate(paths):
spec, original_model, _ = load_entity_source(
entity_spec,
source_path,
mirror=entity_spec.mirror_of is not None,
)
model_file = str(Path(directory.name) / f"{entity_spec.name}-{variant}.xml")
spec.to_file(model_file)
item = scan_genesis_model_metadata(mujoco, SceneCfg(model_file=model_file))
item = _genesis_portable_source_defaults(
mujoco,
owner,
original_model,
item,
scene.default_keyframe_name,
)
_validate_portable_source(owner, item)
files.append(model_file)
metadata.append(item)
entities.append(
GenesisPortableEntitySource(entity_spec.name, tuple(files), tuple(metadata))
)
except BaseException:
directory.cleanup()
raise
return GenesisPortableSources(tuple(entities), directory)
def _genesis_portable_source_defaults(
mujoco: Any,
owner: Any,
source_model: Any,
normalized_metadata: GenesisModelMetadata,
keyframe_name: str | None,
) -> GenesisModelMetadata:
"""Freeze common portable default-keyframe semantics for one source.
The original source model is the only place where MuJoCo-resolved key
qvel/ctrl values remain available. As in the common compiler, only scalar
joint state and actuator controls are overlaid; source root pose/velocity
never replace the portable declaration.
"""
qpos = normalized_metadata.default_qpos.copy()
qvel = np.zeros(normalized_metadata.nv, dtype=np.float32)
ctrl = np.zeros(len(normalized_metadata.actuator_names), dtype=np.float32)
if keyframe_name is None:
return replace(normalized_metadata, default_qvel=qvel, default_ctrl=ctrl)
key_id = int(mujoco.mj_name2id(source_model, mujoco.mjtObj.mjOBJ_KEY, keyframe_name))
if key_id < 0:
return replace(normalized_metadata, default_qvel=qvel, default_ctrl=ctrl)
for field in (source_model.key_qpos, source_model.key_qvel, source_model.key_ctrl):
if not np.isfinite(np.asarray(field[key_id])).all():
raise ValueError(
f"genesis entity {owner.name!r}: keyframe {keyframe_name!r} must be finite"
)
normalized_qpos_addresses = dict(
zip(
normalized_metadata.joint_names,
normalized_metadata.joint_qpos_adrs,
strict=True,
)
)
normalized_qvel_addresses = dict(
zip(
normalized_metadata.joint_names,
normalized_metadata.joint_dof_adrs,
strict=True,
)
)
for joint in owner.joints:
source_id = int(mujoco.mj_name2id(source_model, mujoco.mjtObj.mjOBJ_JOINT, joint.name))
if source_id < 0:
raise ValueError(
f"genesis entity {owner.name!r}: keyframe source lost joint {joint.name}"
)
source_qpos = int(source_model.jnt_qposadr[source_id])
source_qvel = int(source_model.jnt_dofadr[source_id])
qpos[normalized_qpos_addresses[joint.name]] = source_model.key_qpos[key_id, source_qpos]
qvel[normalized_qvel_addresses[joint.name]] = source_model.key_qvel[key_id, source_qvel]
for actuator_name, target_id in zip(owner.actuator_names, owner.actuator_indices, strict=True):
source_id = int(
mujoco.mj_name2id(source_model, mujoco.mjtObj.mjOBJ_ACTUATOR, actuator_name)
)
if source_id < 0:
raise ValueError(
f"genesis entity {owner.name!r}: keyframe source lost actuator {actuator_name}"
)
ctrl[int(target_id)] = source_model.key_ctrl[key_id, source_id]
if not (np.isfinite(qpos).all() and np.isfinite(qvel).all() and np.isfinite(ctrl).all()):
raise RuntimeError(
f"genesis entity {owner.name!r}: normalized default keyframe is nonfinite"
)
return replace(
normalized_metadata,
default_qpos=qpos,
default_qvel=qvel,
default_ctrl=ctrl,
)
def _validate_portable_source(owner: Any, metadata: GenesisModelMetadata) -> None:
"""Compare one normalized source with the frozen public entity layout."""
expected_joints = tuple(joint.name for joint in owner.joints)
if metadata.joint_names != expected_joints:
raise RuntimeError(
f"genesis entity {owner.name!r} joint names/order {metadata.joint_names} "
f"differ from the public layout {expected_joints}"
)
if metadata.nq != len(owner.qpos_indices) or metadata.nv != len(owner.qvel_indices):
raise RuntimeError(
f"genesis entity {owner.name!r} state dimensions "
f"{metadata.nq}/{metadata.nv} differ from the public layout "
f"{len(owner.qpos_indices)}/{len(owner.qvel_indices)}"
)
expected_root = (7, 6) if owner.root_mode == "floating" else (0, 0)
if (metadata.root_qpos_dim, metadata.root_qvel_dim) != expected_root:
raise RuntimeError(
f"genesis entity {owner.name!r} root dimensions differ from the public layout"
)
if metadata.actuator_names != tuple(owner.actuator_names):
raise RuntimeError(
f"genesis entity {owner.name!r} actuator names/order differ from the public layout"
)
if metadata.actuator_joint_names != tuple(owner.actuator_joint_names):
raise RuntimeError(
f"genesis entity {owner.name!r} actuator targets differ from the public layout"
)
if metadata.geom_names != tuple(geom.name for geom in owner.geoms) or (
metadata.geom_body_names != tuple(geom.body_name for geom in owner.geoms)
):
raise RuntimeError(
f"genesis entity {owner.name!r} geom names/ownership differ from the public layout"
)
[docs]
def validate_genesis_portable_sensor_plans(
mujoco: Any,
sources: GenesisPortableSources,
layout: Any,
composed_plans: tuple[GenesisSensorPlan, ...],
) -> tuple[GenesisSensorPlan, ...]:
"""Bind the bounded portable source and fragment sensor subsets.
Genesis does not import MJCF sensors. Portable support is deliberately
limited to entity-owned site ``FramePos``, ``FrameQuat``, ``Gyro``,
``Velocimeter`` and identity-orientation ``Accelerometer`` declarations
whose complete semantic identity is identical in every source variant. The
returned plans use qualified public sensor/body names; the backend computes
pose/motion values from audited native link state and accelerometer values
from a clean public native IMU.
Scene-level fragments are limited further to world-referenced qualified-site
and qualified-body ``FramePos``/``FrameQuat`` declarations, world-referenced
qualified-site/body ``FrameLinVel``/``FrameAngVel`` declarations, and exact
cross-entity geom-pair ``found``/``netforce`` contact declarations. The
common compiler appends them after all entity-owned source sensors, so the
composed prefix must match the independently audited source plans exactly.
"""
source_by_entity = {source.name: source for source in sources.entities}
public_plans: list[GenesisSensorPlan] = []
public_names: set[str] = set()
for owner in layout.entities:
source = source_by_entity[owner.name]
reference = source.metadata[0].sensor_plans
for variant, metadata in enumerate(source.metadata[1:], start=1):
if metadata.sensor_plans != reference:
raise NotImplementedError(
f"genesis entity {owner.name!r} portable sensor identity differs "
f"between variants 0 and {variant}"
)
for plan in reference:
if plan.object_kind != "site":
raise NotImplementedError(
"genesis portable entity source sensors support site objects only"
)
if plan.kind not in (
"framepos",
"framequat",
"gyro",
"velocimeter",
"accelerometer",
):
raise NotImplementedError(
"genesis portable entity source sensors support only site "
"FramePos/FrameQuat/Gyro/Velocimeter/Accelerometer sensors"
)
expected_dim = 4 if plan.kind == "framequat" else 3
if plan.dim != expected_dim:
raise RuntimeError("genesis portable site sensor dimension disagrees with its type")
if plan.reference_type != int(mujoco.mjtObj.mjOBJ_UNKNOWN) or (plan.reference_id != -1):
raise NotImplementedError(
"genesis portable entity source sensors support only "
"unreferenced site FramePos/FrameQuat/Gyro/Velocimeter/"
"Accelerometer sensors"
)
if plan.body_name not in owner.body_names or not plan.object_name:
raise NotImplementedError(
f"genesis portable site sensor {plan.name!r} must reference a named "
f"site owned by entity {owner.name!r}"
)
if plan.site_pos is None or plan.site_quat is None:
raise RuntimeError("genesis portable site sensor has malformed site identity")
if plan.kind == "accelerometer" and not np.allclose(
plan.site_quat, (1.0, 0.0, 0.0, 0.0)
):
raise NotImplementedError(
f"genesis portable site accelerometer {plan.name!r} requires "
"an identity site orientation; Genesis' rotated IMU offset lane "
"is not validated"
)
public_name = f"{owner.name}/{plan.name}"
if public_name in public_names:
raise RuntimeError(
f"genesis portable site sensor name {public_name!r} is not unique"
)
public_names.add(public_name)
public_plans.append(
replace(
plan,
name=public_name,
body_name=f"{owner.name}/{plan.body_name}",
object_name=f"{owner.name}/{plan.object_name}",
)
)
source_plans = tuple(public_plans)
if composed_plans[: len(source_plans)] != source_plans:
raise RuntimeError(
"genesis portable composed source sensors differ from independent entity sources"
)
fragment_plans = composed_plans[len(source_plans) :]
if len({plan.name for plan in composed_plans}) != len(composed_plans):
raise RuntimeError("genesis portable scene sensor names are not unique")
for plan in fragment_plans:
if plan.kind == "contact":
expected_dim = 3 if plan.contact_netforce else 1
if "/" in plan.name or plan.dim != expected_dim:
raise NotImplementedError(
"genesis portable contact fragments support only unprefixed "
"found/netforce sensors with their authored dimension"
)
if (
plan.contact_geom1_name is None
or plan.contact_geom2_name is None
or plan.contact_geom1_name.count("/") != 1
or plan.contact_geom2_name.count("/") != 1
or plan.contact_geom1_name == plan.contact_geom2_name
):
raise NotImplementedError(
f"genesis portable contact fragment sensor {plan.name!r} must "
"reference two distinct qualified public geoms"
)
public_geoms = {
f"{entity.name}/{geom.name}": (entity.name, geom.body_name)
for entity in layout.entities
for geom in entity.geoms
}
geom1 = public_geoms.get(plan.contact_geom1_name)
geom2 = public_geoms.get(plan.contact_geom2_name)
if geom1 is None or geom2 is None or geom1[0] == geom2[0]:
raise NotImplementedError(
f"genesis portable contact fragment sensor {plan.name!r} must "
"reference two distinct public entities"
)
expected_body = f"{geom1[0]}/{geom1[1]}"
if plan.body_name != expected_body:
raise RuntimeError(
f"genesis portable contact fragment sensor {plan.name!r} has "
"malformed geom/body ownership"
)
continue
if plan.kind not in ("framepos", "framequat", "framelinvel", "frameangvel"):
raise NotImplementedError(
"genesis portable sensor fragments support only world-referenced "
"qualified-site FramePos/FrameQuat sensors, qualified-body "
"FramePos/FrameQuat/FrameLinVel/FrameAngVel sensors, qualified-site "
"FrameLinVel/FrameAngVel sensors or "
"exact geom-pair found/netforce contact sensors"
)
expected_dim = 4 if plan.kind == "framequat" else 3
if plan.dim != expected_dim:
raise RuntimeError("genesis portable frame fragment sensor dimension is malformed")
if plan.reference_type != int(mujoco.mjtObj.mjOBJ_UNKNOWN) or (plan.reference_id != -1):
raise NotImplementedError(
"genesis portable sensor fragments support only world-referenced sensors"
)
entity_name, separator, _ = plan.body_name.partition("/")
owner = next((item for item in layout.entities if item.name == entity_name), None)
public_body_names = (
() if owner is None else tuple(f"{owner.name}/{name}" for name in owner.body_names)
)
if not separator or plan.body_name not in public_body_names:
raise NotImplementedError(
f"genesis portable site fragment sensor {plan.name!r} must reference "
"a qualified public site owner/body"
)
if plan.object_kind == "body":
if "/" in plan.name or plan.body_name != plan.object_name:
raise NotImplementedError(
f"genesis portable body fragment sensor {plan.name!r} must "
"reference its exact qualified public body"
)
if plan.site_pos is not None or plan.site_quat is not None:
raise RuntimeError("genesis portable body fragment sensor has malformed identity")
continue
if plan.object_kind != "site":
raise NotImplementedError(
f"genesis portable sensor fragment {plan.name!r} uses an unsupported object type"
)
if plan.kind not in (
"framepos",
"framequat",
"framelinvel",
"frameangvel",
):
raise NotImplementedError(
f"genesis portable site fragment sensor {plan.name!r} supports only "
"world-referenced FramePos/FrameQuat/FrameLinVel/FrameAngVel forms"
)
site_entity, site_separator, _ = plan.object_name.partition("/")
if (
"/" in plan.name
or not site_separator
or site_entity != entity_name
or plan.object_name.count("/") != 1
):
raise NotImplementedError(
f"genesis portable site fragment sensor {plan.name!r} must reference "
"a uniquely qualified public site"
)
if plan.site_pos is None or plan.site_quat is None:
raise RuntimeError("genesis portable site fragment sensor has malformed identity")
return composed_plans
[docs]
@contextmanager
def preserve_torch_globals(torch: Any) -> Iterator[None]:
"""Snapshot torch global defaults and restore them after ``gs.init``.
``gs.init`` mutates the host process: it forces
``torch.set_default_device("cuda:0")`` on the GPU lane, and passing
``seed=`` would reseed the global RNG (REPORT #1372 §3.5 [10]). The
adapter never passes a seed and restores default device/dtype plus the
torch RNG state so the training process observes no pollution
(REPORT #1372 §5.8).
"""
default_device = torch.get_default_device()
default_dtype = torch.get_default_dtype()
cpu_rng_state = torch.get_rng_state()
cuda_rng_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
try:
yield
finally:
if default_device.type == "cpu" and default_device.index is None:
# ``set_default_device`` installs a process-global DeviceContext even
# for CPU, imposing Torch-function dispatch on every later API call.
# CPU is PyTorch's no-context default, so explicitly clear it here.
torch.set_default_device(None)
else:
torch.set_default_device(default_device)
torch.set_default_dtype(default_dtype)
torch.set_rng_state(cpu_rng_state)
if cuda_rng_state is not None:
torch.cuda.set_rng_state_all(cuda_rng_state)
_SESSION_ACTIVE = False
_SESSION_DESTROYED = False
_SESSION_DEVICE_ID: int | None = None
def _pin_cuda_visible_devices(torch: Any, device_id: int) -> int:
"""Pin ``CUDA_VISIBLE_DEVICES`` to the requested device; return the in-process index.
Quadrants (the Genesis 1.3.x compute runtime) binds its CUDA runtime to
the first entry of ``CUDA_VISIBLE_DEVICES`` regardless of torch's current
device — ``torch.cuda.set_device(N)`` alone leaves the engine on physical
GPU 0 and crashes later with cross-device illegal-memory errors (verified
on genesis_world 1.3.3 / Quadrants 1.3.0, UniLab issue #1508). The only
owner-layer way to place a session on another physical GPU is to shrink
visibility to that GPU before the first CUDA context exists; afterwards
the in-process device index is 0.
``device_id`` addresses the *current* visibility namespace: without the
variable set it is the host index, otherwise it indexes into the existing
entries (physical indices or UUID strings). Returns 0 on success so the
session records the in-process index actually used.
"""
if bool(torch.cuda.is_initialized()):
raise RuntimeError(
f"genesis device_id={device_id} requires pinning CUDA_VISIBLE_DEVICES "
"before any CUDA context is created in this process, but torch CUDA "
"is already initialized; select the device earlier (entrypoint cold "
"path) or set CUDA_VISIBLE_DEVICES before launching the process"
)
raw = os.environ.get("CUDA_VISIBLE_DEVICES")
entries = [entry.strip() for entry in raw.split(",") if entry.strip()] if raw else None
if entries is None:
device_count = int(torch.cuda.device_count())
if device_id >= device_count:
raise ValueError(
f"genesis device_id={device_id} is out of range; "
f"torch.cuda.device_count()={device_count}"
)
target = str(device_id)
else:
if device_id >= len(entries):
raise ValueError(
f"genesis device_id={device_id} is out of range for "
f"CUDA_VISIBLE_DEVICES={raw!r} ({len(entries)} entr(ies))"
)
target = entries[device_id]
os.environ["CUDA_VISIBLE_DEVICES"] = target
# ``device_count`` is lru-cached; drop any pre-pin host-wide count so
# later callers observe the pinned single-device namespace.
cache_clear = getattr(torch.cuda.device_count, "cache_clear", None)
if callable(cache_clear):
cache_clear()
return 0
def _resolve_genesis_device_id(torch: Any, device_id: int | None) -> int | None:
"""Validate and select the CUDA index used by a Genesis session."""
if device_id is not None:
if isinstance(device_id, bool) or not isinstance(device_id, int) or device_id < 0:
raise ValueError(
f"genesis device_id must be a non-negative integer or None, got {device_id!r}"
)
if device_id > 0:
# The engine only honors the first visible device; pin visibility
# and continue with the in-process index 0. The pin must run
# before any torch CUDA query (even ``is_available`` latches
# CUDA_VISIBLE_DEVICES in the runtime), so it precedes the
# availability check below.
device_id = _pin_cuda_visible_devices(torch, device_id)
if not bool(torch.cuda.is_available()):
raise ValueError(
f"genesis device_id={device_id} requires CUDA, but CUDA is unavailable"
)
# Genesis' ``gs.init`` consults the process-wide current CUDA device.
# Bind it before entering the initialization routine; doing this after
# ``gs.init`` is too late because Genesis has already allocated its
# global backend on device zero.
torch.cuda.set_device(device_id)
return int(device_id)
if bool(torch.cuda.is_available()):
return int(torch.cuda.current_device())
return None
[docs]
def init_genesis_session(
deps: Any,
*,
device_id: int | None = None,
logging_level: str = "warning",
) -> None:
"""Initialize the process-wide Genesis session exactly once.
Repeated ``init -> destroy`` cycles are functional but leak 200-450 MB of
host RSS per cycle (REPORT #1372 §3.5 [9a]), so a long-lived training
process gets exactly one session. Multiple GenesisBackend instances share
the live session (multi-scene coexistence is measured OK, REPORT [9b]);
after :func:`destroy_genesis_session` any further construction fails
closed with a clear error.
"""
global _SESSION_ACTIVE, _SESSION_DESTROYED, _SESSION_DEVICE_ID
if _SESSION_DESTROYED:
raise RuntimeError(
"genesis backend supports exactly one gs.init per process and the session "
"was already destroyed; start a fresh process instead of re-initializing "
"(init/destroy cycles leak host RSS, REPORT #1372 §3.5 [9a])."
)
selected_device_id = _resolve_genesis_device_id(deps.torch, device_id)
if _SESSION_ACTIVE:
if (
selected_device_id is not None
and _SESSION_DEVICE_ID is not None
and selected_device_id != _SESSION_DEVICE_ID
):
raise RuntimeError(
"genesis session is already initialized on CUDA device "
f"{_SESSION_DEVICE_ID}, cannot reuse it on device {selected_device_id}"
)
return
gs = deps.genesis
cuda_available = bool(deps.torch.cuda.is_available())
backend_kind = gs.gpu if cuda_available else gs.cpu
try:
with preserve_torch_globals(deps.torch):
gs.init(backend=backend_kind, logging_level=logging_level)
except Exception as exc:
raise RuntimeError(
"genesis backend failed to initialize the Genesis runtime "
f"(backend={'gpu' if cuda_available else 'cpu'}; only the gs.gpu lane is "
f"validated by REPORT #1372): {type(exc).__name__}: {exc}"
) from exc
_SESSION_ACTIVE = True
_SESSION_DEVICE_ID = selected_device_id
[docs]
def destroy_genesis_session(deps: Any) -> None:
"""Tear down the process-wide session; re-initialization stays forbidden."""
global _SESSION_ACTIVE, _SESSION_DESTROYED, _SESSION_DEVICE_ID
if not _SESSION_ACTIVE:
return
deps.genesis.destroy()
_SESSION_ACTIVE = False
_SESSION_DESTROYED = True
_SESSION_DEVICE_ID = None
def _reset_session_state_for_tests() -> None:
"""Reset the lifecycle guard; test-only seam for the fake runtime lane."""
global _SESSION_ACTIVE, _SESSION_DESTROYED, _SESSION_DEVICE_ID
_SESSION_ACTIVE = False
_SESSION_DESTROYED = False
_SESSION_DEVICE_ID = None
def _map_global_option(name: str, value: str, table: dict[str, str], enum_ns: Any) -> Any:
try:
attr = table[value]
except KeyError as exc:
supported = ", ".join(sorted(table))
raise ValueError(
f"{name} must be one of: {supported}; got {value!r}. Genesis drops the MJCF "
"global <option> block, so the owner YAML must declare a supported value."
) from exc
return getattr(enum_ns, attr)
[docs]
def build_scene(
deps: Any,
*,
sim_dt: float,
gravity: np.ndarray,
integrator: str | None,
constraint_solver: str | None,
friction_cone: str | None,
solver_iterations: int | None,
) -> Any:
"""Construct the unbuilt Genesis scene with explicit global options."""
gs = deps.genesis
rigid_kwargs: dict[str, Any] = {
# Per-env DR setters (mass/frictionloss/kp/kv) are materialize-time
# decisions and require batched link/dof info (REPORT #1372 §5.7).
"batch_links_info": True,
"batch_dofs_info": True,
}
if integrator is not None:
rigid_kwargs["integrator"] = _map_global_option(
"genesis_integrator", integrator, _INTEGRATOR_ENUMS, gs.integrator
)
if constraint_solver is not None:
rigid_kwargs["constraint_solver"] = _map_global_option(
"genesis_constraint_solver",
constraint_solver,
_CONSTRAINT_SOLVER_ENUMS,
gs.constraint_solver,
)
if friction_cone is not None:
rigid_kwargs["friction_cone"] = _map_global_option(
"genesis_friction_cone", friction_cone, _FRICTION_CONE_ENUMS, gs.friction_cone
)
if solver_iterations is not None:
rigid_kwargs["iterations"] = int(solver_iterations)
gravity_tuple = tuple(float(component) for component in np.asarray(gravity, dtype=np.float64))
return gs.Scene(
sim_options=gs.options.SimOptions(dt=float(sim_dt), gravity=gravity_tuple),
rigid_options=gs.options.RigidOptions(**rigid_kwargs),
show_viewer=False,
)
def _scan_sensor_plans(
mujoco: Any, model: Any, *, allow_cross_entity_contacts: bool = False
) -> tuple[GenesisSensorPlan, ...]:
"""Map the MJCF sensor table onto Genesis equivalents (REPORT §5.3).
Genesis 1.3.3 does not import ``<sensor>`` at all. Supported mappings:
gyro/accelerometer/velocimeter -> IMU-class site sensors computed from
link state; framepos/framequat/framezaxis -> link state plus site-frame
math; contact ``data="found"`` -> per-link net contact force threshold.
Portable fragment contact plans additionally accept the compiler's exact
geom-pair ``force/netforce`` identity. Anything else fails closed here, at
the nearest cold path.
"""
sensor_obj = mujoco.mjtObj.mjOBJ_SENSOR
plans: list[GenesisSensorPlan] = []
for sensor_id in range(int(model.nsensor)):
name = mujoco.mj_id2name(model, sensor_obj, sensor_id)
if not name:
raise NotImplementedError(
f"genesis backend requires named MJCF sensors; sensor id {sensor_id} is unnamed"
)
sensor_type = mujoco.mjtSensor(int(model.sensor_type[sensor_id]))
dim = int(model.sensor_dim[sensor_id])
if sensor_type == mujoco.mjtSensor.mjSENS_CONTACT:
plans.append(
_scan_contact_sensor(
mujoco,
model,
sensor_id,
str(name),
dim,
allow_cross_entity_contacts=allow_cross_entity_contacts,
)
)
continue
site_kinds = {
mujoco.mjtSensor.mjSENS_GYRO: "gyro",
mujoco.mjtSensor.mjSENS_ACCELEROMETER: "accelerometer",
mujoco.mjtSensor.mjSENS_VELOCIMETER: "velocimeter",
mujoco.mjtSensor.mjSENS_FRAMEPOS: "framepos",
mujoco.mjtSensor.mjSENS_FRAMEQUAT: "framequat",
mujoco.mjtSensor.mjSENS_FRAMELINVEL: "framelinvel",
mujoco.mjtSensor.mjSENS_FRAMEANGVEL: "frameangvel",
mujoco.mjtSensor.mjSENS_FRAMEZAXIS: "framezaxis",
}
kind = site_kinds.get(sensor_type)
if kind is None:
raise NotImplementedError(
f"genesis backend cannot map MJCF sensor {name!r} of type "
f"{sensor_type.name}; supported types: contact(found), gyro, accelerometer, "
"velocimeter, framepos, framequat, framezaxis (REPORT #1372 §3.4)."
)
sensor_objtype = int(model.sensor_objtype[sensor_id])
composed_body = (
allow_cross_entity_contacts
and kind in ("framepos", "framequat", "framelinvel", "frameangvel")
and sensor_objtype == int(mujoco.mjtObj.mjOBJ_BODY)
)
if composed_body:
body_id = int(model.sensor_objid[sensor_id])
body_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id)
if not body_name:
raise NotImplementedError(
f"genesis body sensor {name!r} references unnamed body id {body_id}"
)
reference_type = int(model.sensor_reftype[sensor_id])
reference_id = int(model.sensor_refid[sensor_id])
if reference_type != int(mujoco.mjtObj.mjOBJ_UNKNOWN) or reference_id != -1:
raise NotImplementedError(
f"genesis backend maps {kind} sensor {name!r} only with a world reference"
)
plans.append(
GenesisSensorPlan(
name=str(name),
kind=kind,
dim=dim,
body_name=str(body_name),
object_name=str(body_name),
reference_type=reference_type,
reference_id=reference_id,
site_pos=None,
site_quat=None,
object_kind="body",
)
)
continue
if sensor_objtype != int(mujoco.mjtObj.mjOBJ_SITE):
raise NotImplementedError(
f"genesis backend maps {kind} sensors from MJCF sites only; sensor {name!r} "
f"uses objtype {sensor_objtype}."
)
site_id = int(model.sensor_objid[sensor_id])
body_id = int(model.site_bodyid[site_id])
site_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_SITE, site_id)
if not site_name:
raise NotImplementedError(
f"genesis site sensor {name!r} references unnamed site id {site_id}"
)
body_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id)
site_pos = tuple(float(v) for v in np.asarray(model.site_pos[site_id], dtype=np.float64))
site_quat = tuple(float(v) for v in np.asarray(model.site_quat[site_id], dtype=np.float64))
reference_type = int(model.sensor_reftype[sensor_id])
reference_id = int(model.sensor_refid[sensor_id])
if kind in ("framepos", "framequat", "framezaxis", "framelinvel", "frameangvel") and (
reference_type != int(mujoco.mjtObj.mjOBJ_UNKNOWN) or reference_id != -1
):
raise NotImplementedError(
f"genesis backend maps {kind} sensor {name!r} only with a world reference"
)
if kind == "accelerometer" and not np.allclose(site_quat, (1.0, 0.0, 0.0, 0.0)):
raise NotImplementedError(
f"genesis backend maps accelerometer {name!r} onto an IMUSensor, whose "
"euler_offset path is not a validated support lane; rotated accelerometer "
"sites are rejected (gyro/velocimeter/frame sensors support rotation)."
)
plans.append(
GenesisSensorPlan(
name=str(name),
kind=kind,
dim=dim,
body_name=str(body_name),
object_name=str(site_name),
reference_type=reference_type,
reference_id=reference_id,
site_pos=site_pos,
site_quat=site_quat,
)
)
return tuple(plans)
def _scan_contact_sensor(
mujoco: Any,
model: Any,
sensor_id: int,
name: str,
dim: int,
*,
allow_cross_entity_contacts: bool = False,
) -> GenesisSensorPlan:
intprm = tuple(int(value) for value in np.asarray(model.sensor_intprm[sensor_id])[:3])
contact_netforce = False
# (data, reduce, num): the found flag ignores the reduce mode entirely
# (num=1 leaves nothing to aggregate), so none (0) and mindist (1) compile
# to the same observable scalar; accept both.
if intprm in ((1, 0, 1), (1, 1, 1)):
expected_dim = 1
elif allow_cross_entity_contacts and intprm == (2, 3, 1):
contact_netforce = True
expected_dim = 3
else:
raise NotImplementedError(
"genesis backend maps only the compiled geom-pair found and portable "
"cross-entity netforce contact forms; "
f"sensor {name!r} has compiled contact parameters {intprm}."
)
if dim != expected_dim:
raise NotImplementedError(
f"genesis contact sensor {name!r} dimension {dim} disagrees with its "
f"compiled form (expected {expected_dim})."
)
if int(model.sensor_objtype[sensor_id]) != int(mujoco.mjtObj.mjOBJ_GEOM) or int(
model.sensor_reftype[sensor_id]
) != int(mujoco.mjtObj.mjOBJ_GEOM):
raise NotImplementedError(
f"genesis backend maps contact sensors over geom pairs only; sensor {name!r} "
"uses a non-geom object type."
)
geom1_id = int(model.sensor_objid[sensor_id])
geom2_id = int(model.sensor_refid[sensor_id])
body1_id = int(model.geom_bodyid[geom1_id])
body2_id = int(model.geom_bodyid[geom2_id])
if not allow_cross_entity_contacts and (body1_id > 0) == (body2_id > 0):
raise NotImplementedError(
f"genesis backend maps contact sensor {name!r} onto the robot-side link's net "
"contact force; exactly one geom must belong to the world body "
f"(got body ids {body1_id}/{body2_id})."
)
body_id = body1_id if body1_id > 0 else body2_id
body_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id)
geom1_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, geom1_id)
geom2_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, geom2_id)
if not geom1_name:
raise NotImplementedError(
f"genesis contact sensor {name!r} references unnamed geom id {geom1_id}"
)
if not geom2_name:
raise NotImplementedError(
f"genesis contact sensor {name!r} references unnamed geom id {geom2_id}"
)
return GenesisSensorPlan(
name=name,
kind="contact",
dim=dim,
body_name=str(body_name),
object_name=str(geom1_name),
reference_type=int(model.sensor_reftype[sensor_id]),
reference_id=geom2_id,
site_pos=None,
site_quat=None,
contact_geom1_name=str(geom1_name),
contact_geom2_name=str(geom2_name),
contact_netforce=contact_netforce,
)