"""Cold-path scene materialization for the independent ``genesis`` backend.
Genesis 1.3.3 drops three MJCF features UniLab relies on (measured in
``scripts/tools/genesis_feasibility/REPORT.md`` #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
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Iterator
import numpy as np
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
class _TemporarySceneCleanup:
"""Own the temporary XMLs created while materializing one scene."""
def __init__(self, *paths: str) -> None:
self._paths = paths
self._cleaned = False
def cleanup(self) -> None:
if self._cleaned:
return
self._cleaned = True
for path in self._paths:
try:
os.remove(path)
except FileNotFoundError:
pass
[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 sensors.
"""
name: str
kind: str
dim: int
body_name: str
site_pos: tuple[float, ...] | None
site_quat: tuple[float, ...] | None
[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:
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_genesis_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) -> 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.
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))
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_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)."
)
if int(model.sensor_objtype[sensor_id]) != int(mujoco.mjtObj.mjOBJ_SITE):
raise NotImplementedError(
f"genesis backend maps {kind} sensors from MJCF sites only; sensor {name!r} "
f"uses objtype {int(model.sensor_objtype[sensor_id])}."
)
site_id = int(model.sensor_objid[sensor_id])
body_id = int(model.site_bodyid[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))
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),
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
) -> GenesisSensorPlan:
if dim != 1:
raise NotImplementedError(
f'genesis backend maps contact sensors with data="found" only (dim 1); '
f"sensor {name!r} has dim {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 (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)
return GenesisSensorPlan(
name=name,
kind="contact",
dim=dim,
body_name=str(body_name),
site_pos=None,
site_quat=None,
)