import logging
import os
import time
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, TypeVar, cast
import numpy as np
from unisim.dr.types import (
INTERVAL_TERM_BODY_FORCE,
INTERVAL_TERM_PUSH,
RESET_TERM_BASE_COM,
RESET_TERM_BASE_MASS,
RESET_TERM_BODY_IPOS,
RESET_TERM_BODY_MASS,
RESET_TERM_GEOM_FRICTION,
RESET_TERM_GRAVITY,
RESET_TERM_KD,
RESET_TERM_KP,
DomainRandomizationCapabilities,
InitRandomizationPlan,
IntervalTermOp,
ResetRandomizationPayload,
)
from unisim.scene import SceneCfg
from unisim.utils.rotation import np_quat_apply_inverse_batched
try:
import motrixsim as mtx
from motrixsim.render import RenderApp, RenderSettings
from motrixsim.render import RenderClosedError as _MotrixRenderClosedError
MOTRIX_AVAILABLE = True
except ImportError:
MOTRIX_AVAILABLE = False
# No motrixsim in this process: the ``except _MotrixRenderClosedError``
# clauses below never match, which is correct because the renderer cannot
# exist without the package.
_MotrixRenderClosedError = ()
from ..base import (
BackendHeightScanner,
BackendPlayCapabilities,
BackendPlayRenderPlan,
BackendRootStateLayout,
BackendTerrainSpawnData,
CameraCfg,
RenderClosedError,
SimBackend,
normalize_play_render_mode,
unsupported_debug_overlay_error,
)
from ..motrix_camera import (
MotrixTrackingCamera,
render_offsets,
resolve_system_camera_view,
tracking_camera_lookat,
)
from .playback import run_motrix_playback
logger = logging.getLogger(__name__)
T = TypeVar("T")
DEFAULT_MOTRIX_MAX_ITERATIONS = 3
def _require_not_none(value: T | None, error_message: str) -> T:
if value is None:
raise ValueError(error_message)
return value
def _first_scalar(value: Any) -> float:
arr = np.asarray(value, dtype=np.float32)
return float(arr.reshape(-1)[0])
def _contiguous_slice(indices: np.ndarray) -> slice | None:
if indices.size == 0:
return None
start = int(indices[0])
stop = start + int(indices.size)
if np.array_equal(indices, np.arange(start, stop, dtype=indices.dtype)):
return slice(start, stop)
return None
@dataclass
class _MotrixSceneContext:
model: "mtx.SceneModel"
sensor_names: tuple[str, ...]
terrain_origins: np.ndarray | None = None
terrain_surface_sampler: object | None = None
cleanup_handle: object | None = None
@dataclass
class _MotrixTerrainScanner(BackendHeightScanner):
scanner: "mtx.TerrainScanner"
data: "mtx.SceneData"
out: np.ndarray
def scan(self) -> np.ndarray:
heights = np.asarray(self.scanner.scan(self.data, out=self.out))
if heights.shape != self.out.shape:
raise ValueError(
f"Motrix TerrainScanner.scan returned shape {heights.shape}, "
f"expected {self.out.shape}"
)
return heights
def _build_motrix_scene_context(
scene: SceneCfg,
*,
add_body_sensors: bool,
base_name: str,
) -> _MotrixSceneContext:
from unisim.backend.motrix.scene import (
_materialize_motrix_hfield_attached_scene_with_sensor_names,
_materialize_motrix_scene_with_sensor_names,
)
if scene is None:
raise ValueError("SceneCfg must be provided")
if not scene.model_file:
raise ValueError("SceneCfg.model_file must be provided")
if scene.terrain is None:
model, sensor_names = _materialize_motrix_scene_with_sensor_names(
model_file=scene.model_file,
fragment_files=scene.fragment_files,
add_body_sensors=add_body_sensors,
base_name=base_name,
)
return _MotrixSceneContext(model=model, sensor_names=sensor_names)
if scene.terrain.generator is None:
raise ValueError("SceneCfg.terrain.generator must be configured for terrain scenes")
model, terrain_origins, terrain_surface_sampler, sensor_names = (
_materialize_motrix_hfield_attached_scene_with_sensor_names(
model_file=scene.model_file,
terrain_cfg=scene.terrain.generator,
fragment_files=scene.fragment_files,
hfield_name=scene.terrain.hfield_name,
geom_name=scene.terrain.geom_name or "floor",
add_body_sensors=add_body_sensors,
base_name=base_name,
return_surface_sampler=True,
)
)
return _MotrixSceneContext(
model=model,
sensor_names=sensor_names,
terrain_origins=terrain_origins,
terrain_surface_sampler=terrain_surface_sampler,
)
[docs]
class MotrixBackend(SimBackend):
"""MotrixSim backend implementation."""
[docs]
def __init__(
self,
scene: SceneCfg,
num_envs: int,
sim_dt: float,
base_name: str = "base",
np_dtype=np.float32,
add_body_sensors: bool = False,
max_iterations: int | None = DEFAULT_MOTRIX_MAX_ITERATIONS,
push_body_name: str | None = None,
):
if not MOTRIX_AVAILABLE:
raise ImportError("motrixsim not available")
scene_context = _build_motrix_scene_context(
scene,
add_body_sensors=add_body_sensors,
base_name=base_name,
)
self._scene = scene
self.scene_artifacts_dir = None
self.terrain_origins = scene_context.terrain_origins
self.terrain_surface_sampler = scene_context.terrain_surface_sampler
self._terrain_spawn_data = (
None
if self.terrain_origins is None
else BackendTerrainSpawnData(
terrain_origins=self.terrain_origins,
sample_height=(
None
if self.terrain_surface_sampler is None
else cast(Any, self.terrain_surface_sampler).sample_height
),
)
)
self._scene_cleanup_handle = scene_context.cleanup_handle
self._base_name = base_name
self._model = scene_context.model
self._sensor_names = frozenset(scene_context.sensor_names)
self._body_id_to_name = { # type: ignore[assignment]
link.index: link.name for link in self._model.links if link.name
}
self._model.options.timestep = sim_dt
if max_iterations is None:
max_iterations = DEFAULT_MOTRIX_MAX_ITERATIONS
self._model.options.max_iterations = int(max_iterations)
self._num_envs = num_envs
self._np_dtype = np_dtype
self._pre_step_control_fn = None
self._data = mtx.SceneData(self._model, batch=[num_envs]) # pyright: ignore[reportPossiblyUnbound]
self._body: "mtx.Body" = _require_not_none(
self._model.get_body(base_name), f"Body '{base_name}' not found in Motrix model"
)
self._body_link: "mtx.Link" = _require_not_none(
self._model.get_link(base_name), f"Link '{base_name}' not found in Motrix model"
)
push_body = push_body_name if push_body_name is not None else base_name
self._push_body_link: "mtx.Link" = _require_not_none(
self._model.get_link(push_body), f"Push link '{push_body}' not found in Motrix model"
)
self._body_floatingbase = self._body.floatingbase
self._joint_dof_pos_indices = np.asarray(self._model.joint_dof_pos_indices, dtype=np.intp)
self._joint_dof_vel_indices = np.asarray(self._model.joint_dof_vel_indices, dtype=np.intp)
self._joint_dof_pos_slice = _contiguous_slice(self._joint_dof_pos_indices)
position_actuators: list["mtx.PositionActuator"] = []
for actuator in self._model.actuators:
if actuator.typ == "position":
position_actuators.append(cast("mtx.PositionActuator", actuator))
self._position_actuators = position_actuators
self._supports_position_actuator_gains = len(self._position_actuators) == int(
self._model.num_actuators
)
# qpos index of each position actuator's own target joint, in actuator
# order. Used to reset position actuators to "hold current pose" without
# assuming a fully-actuated model: parallel / under-actuated mechanisms
# (e.g. a Stewart platform) have passive joints, so the model-wide
# ``joint_dof_pos_indices`` is wider than ``num_actuators``.
self._actuator_joint_pos_indices: np.ndarray | None = None
self._actuator_joint_vel_indices: np.ndarray | None = None
if self._supports_position_actuator_gains:
joint_pos_idx: list[int] = []
joint_vel_idx: list[int] = []
for actuator in sorted(self._position_actuators, key=lambda a: int(a.index)):
if actuator.target_type != "joint":
joint_pos_idx = []
joint_vel_idx = []
break
joint = self._model.get_joint(actuator.target_name)
if joint is None or int(joint.num_dof_pos) != 1:
joint_pos_idx = []
joint_vel_idx = []
break
joint_pos_idx.append(int(joint.dof_pos_index))
joint_vel_idx.append(int(joint.dof_vel_index))
if len(joint_pos_idx) == int(self._model.num_actuators):
self._actuator_joint_pos_indices = np.asarray(joint_pos_idx, dtype=np.intp)
self._actuator_joint_vel_indices = np.asarray(joint_vel_idx, dtype=np.intp)
self._actuator_joint_pos_slice = (
_contiguous_slice(self._actuator_joint_pos_indices)
if self._actuator_joint_pos_indices is not None
else None
)
self._default_actuator_kp = np.zeros((self.num_actuators,), dtype=np.float32)
self._default_actuator_kd = np.zeros((self.num_actuators,), dtype=np.float32)
for actuator in self._position_actuators:
idx = int(actuator.index)
# TODO: switch to motrixsim model-level actuator gain API once available.
self._default_actuator_kp[idx] = _first_scalar(actuator.get_kp_override(self._data))
self._default_actuator_kd[idx] = _first_scalar(actuator.get_kd_override(self._data))
self._floating_base_quat_indices: tuple[np.ndarray, ...] = tuple(
np.asarray(floating_base.dof_pos_indices[3:7], dtype=np.intp)
for floating_base in getattr(self._model, "floating_bases", [])
if len(floating_base.dof_pos_indices) >= 7
)
self._links_by_id: dict[int, "mtx.Link"] = {
int(link.index): link for link in self._model.links
}
self._supports_external_force = all(
callable(getattr(link, "add_external_force", None))
for link in self._links_by_id.values()
)
self._applied_body_forces: dict[int, np.ndarray] = {}
self._geoms_by_id: dict[int, "mtx.Geom"] = {
int(geom.index): geom for geom in self._model.geoms
}
# TODO(motrixsim): once pure visual geoms either stop exposing friction
# override methods or safely no-op them, drop this collision-mask filter.
self._geom_friction_override_ids = tuple(
geom_id
for geom_id, geom in self._geoms_by_id.items()
if (
int(getattr(geom, "collision_group", 0)) != 0
or int(getattr(geom, "collision_affinity", 0)) != 0
)
)
self._supports_geom_friction_override = all(
callable(getattr(geom, "get_friction_override", None))
and callable(getattr(geom, "set_friction_override", None))
for geom_id, geom in self._geoms_by_id.items()
if geom_id in self._geom_friction_override_ids
)
self._supports_gravity_override = callable(
getattr(self._model, "get_gravity_override", None)
) and callable(getattr(self._model, "set_gravity_override", None))
self._default_body_mass = np.zeros((int(self._model.num_links),), dtype=np.float32)
self._default_body_ipos = np.zeros((int(self._model.num_links), 3), dtype=np.float32)
for link_id, link in self._links_by_id.items():
self._default_body_mass[link_id] = _first_scalar(link.get_mass_override(self._data))
self._default_body_ipos[link_id] = np.asarray(
link.get_center_of_mass_override(self._data),
dtype=np.float32,
).reshape(self._num_envs, 3)[0]
self._default_geom_friction = np.zeros((int(self._model.num_geoms), 3), dtype=np.float32)
if self._supports_geom_friction_override:
for geom_id in self._geom_friction_override_ids:
geom = self._geoms_by_id[geom_id]
self._default_geom_friction[geom_id] = np.asarray(
geom.get_friction_override(self._data),
dtype=np.float32,
).reshape(self._num_envs, 3)[0]
self._init_geom_size_overrides: dict[int, np.ndarray] = {}
self._render_app: "RenderApp | None" = None
self._render_headless: bool | None = None
self._render_capture_enabled = False
self._render_offsets_np: np.ndarray | None = None
self._render_tracking_camera: MotrixTrackingCamera | None = None
self.backend_type = "motrix"
self._link_velocity_cache: np.ndarray | None = None
# Pre-cache link objects to avoid repeated get_link() lookups.
self._link_cache: dict[int, "mtx.Link"] = {}
for link in self._model.links:
if link.name:
self._link_cache[link.index] = link
# Run forward kinematics once so initial link poses and sensor data are valid.
self._model.forward_kinematic(self._data)
self._link_velocities: np.ndarray | None = None
self._link_velocity_cache_valid = False
self._refresh_link_pose_cache()
# Scratch buffers reused by set_state() to avoid per-reset allocations.
# Sized to the full env count and rewritten in place each call.
self._set_state_mask_scratch: np.ndarray = np.zeros(self._num_envs, dtype=bool)
self._set_state_qpos_motrix_scratch: np.ndarray | None = None
[docs]
def get_motion_body_ids(self, names: Sequence[str]) -> np.ndarray:
ids: list[int] = []
for name in names:
link_id = self._model.get_link_index(name)
if link_id is None or link_id < 0:
raise ValueError(f"Motion body '{name}' not found in Motrix model")
# Motion datasets use MuJoCo-style body ids, where worldbody is id 0.
ids.append(int(link_id) + 1)
return np.array(ids, dtype=np.int32)
# ------------------------------------------------------------------ #
# Properties #
# ------------------------------------------------------------------ #
@property
def num_envs(self) -> int:
return self._num_envs
@property
def model(self):
return self._model
@property
def data(self):
return self._data
# ------------------------------------------------------------------ #
# Model properties #
# ------------------------------------------------------------------ #
@property
def num_actuators(self) -> int:
return int(self._model.num_actuators)
@property
def num_dof_vel(self) -> int:
return int(len(self._joint_dof_vel_indices))
[docs]
def get_actuator_ctrl_range(self) -> np.ndarray:
arr: np.ndarray = np.array(self._model.actuator_ctrl_limits, dtype=self._np_dtype)
result: np.ndarray = arr.T.copy()
return result
[docs]
def get_actuator_names(self) -> tuple[str, ...]:
actuators = sorted(self._model.actuators, key=lambda actuator: int(actuator.index))
names = tuple(str(actuator.name) for actuator in actuators)
if len(names) != self.num_actuators or any(not name for name in names):
raise NotImplementedError(
"backend 'motrix' capability 'actuator names' requires one non-empty name "
f"per control column; received {names}"
)
if len(set(names)) != len(names):
raise NotImplementedError(
"backend 'motrix' capability 'actuator names' requires unique names; "
f"received {names}"
)
return names
[docs]
def get_actuator_joint_names(self) -> tuple[str, ...]:
actuators = sorted(self._model.actuators, key=lambda actuator: int(actuator.index))
names: list[str] = []
for actuator in actuators:
if actuator.target_type != "joint":
raise NotImplementedError(
"backend 'motrix' capability 'actuator target joint' requires a joint "
f"transmission; actuator '{actuator.name}' targets '{actuator.target_type}'"
)
joint = self._model.get_joint(actuator.target_name)
if joint is None or int(joint.num_dof_pos) != 1 or int(joint.num_dof_vel) != 1:
raise NotImplementedError(
"backend 'motrix' capability 'actuator target joint' requires a "
f"single-DoF joint; actuator '{actuator.name}' targets "
f"'{actuator.target_name}'"
)
names.append(str(actuator.target_name))
if len(names) != self.num_actuators:
raise NotImplementedError(
"backend 'motrix' capability 'actuator target joint' returned "
f"{len(names)} targets for {self.num_actuators} actuators"
)
return tuple(names)
[docs]
def get_terrain_spawn_data(self) -> BackendTerrainSpawnData | None:
return self._terrain_spawn_data
[docs]
def get_keyframe_qpos(self, name: str) -> np.ndarray:
if hasattr(self._model, "keyframes") and self._model.num_keyframes > 0:
qpos = np.array(self._model.keyframes[0].dof_pos, dtype=self._np_dtype)
else:
qpos = np.array(self._model.compute_init_dof_pos(), dtype=self._np_dtype)
return self._motrix_qpos_to_mujoco(qpos)
[docs]
def get_default_qpos(self) -> np.ndarray:
qpos = np.array(self._model.compute_init_dof_pos(), dtype=self._np_dtype)
return self._motrix_qpos_to_mujoco(qpos)
[docs]
def get_default_dof_pos(self) -> np.ndarray:
qpos = np.asarray(self._model.compute_init_dof_pos(), dtype=self._np_dtype)
indices = (
self._actuator_joint_pos_indices
if self._actuator_joint_pos_indices is not None
else self._joint_dof_pos_indices
)
return np.asarray(qpos[indices], dtype=self._np_dtype).copy()
[docs]
def get_init_qvel(self) -> np.ndarray:
return np.zeros((self._model.num_dof_vel,), dtype=self._np_dtype)
[docs]
def get_root_state_layout(self, root_body_name: str) -> BackendRootStateLayout:
body = self._model.get_body(root_body_name)
if body is None:
raise ValueError(f"Body '{root_body_name}' not found in Motrix model")
floating_base = body.floatingbase
if floating_base is None:
raise NotImplementedError(
"backend 'motrix' capability 'root-state layout' requires body "
f"'{root_body_name}' to own a floating base"
)
return BackendRootStateLayout(
qpos_indices=tuple(int(index) for index in floating_base.dof_pos_indices),
qvel_indices=tuple(int(index) for index in floating_base.dof_vel_indices),
)
[docs]
def get_body_ids(self, names: Sequence[str]) -> np.ndarray:
ids: list[int] = []
for name in names:
bid = self._model.get_link_index(name)
if bid is None or bid < 0:
raise ValueError(f"Body '{name}' not found in Motrix model")
ids.append(int(bid))
return np.array(ids, dtype=np.int32)
[docs]
def get_site_ids(self, names: Sequence[str]) -> np.ndarray:
ids: list[int] = []
for name in names:
sid = self._model.get_site_index(name)
if sid is None or sid < 0:
raise ValueError(f"Site '{name}' not found in Motrix model")
ids.append(int(sid))
return np.array(ids, dtype=np.int32)
[docs]
def get_joint_dof_indices(self, names: Sequence[str]) -> np.ndarray:
indices: list[int] = []
for name in names:
joint = self._resolve_single_dof_joint(name)
indices.append(int(joint.dof_vel_index))
return np.array(indices, dtype=np.int32)
[docs]
def get_joint_dof_pos_indices(self, names: Sequence[str]) -> np.ndarray:
indices: list[int] = []
for name in names:
joint = self._resolve_single_dof_joint(name)
indices.append(self._joint_dof_local_index(name, int(joint.dof_pos_index), pos=True))
return np.array(indices, dtype=np.int32)
[docs]
def get_joint_dof_vel_indices(self, names: Sequence[str]) -> np.ndarray:
indices: list[int] = []
for name in names:
joint = self._resolve_single_dof_joint(name)
indices.append(self._joint_dof_local_index(name, int(joint.dof_vel_index), pos=False))
return np.array(indices, dtype=np.int32)
[docs]
def get_joint_state_qpos_indices(self, names: Sequence[str]) -> np.ndarray:
indices = [int(self._resolve_single_dof_joint(name).dof_pos_index) for name in names]
return np.asarray(indices, dtype=np.int32)
[docs]
def get_joint_state_qvel_indices(self, names: Sequence[str]) -> np.ndarray:
indices = [int(self._resolve_single_dof_joint(name).dof_vel_index) for name in names]
return np.asarray(indices, dtype=np.int32)
[docs]
def get_site_jacobian_w(
self,
site_id: int,
dof_indices: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
sid = int(site_id)
if sid < 0 or sid >= int(self._model.num_sites):
raise ValueError(f"site_id out of range: {sid}")
site = self._model.sites[sid]
jac = np.asarray(site.get_jacobian(self._data), dtype=self._np_dtype)
if jac.ndim != 3 or jac.shape[0] != self._num_envs or jac.shape[1] != 6:
raise ValueError(
f"Motrix site Jacobian for site {sid} must have shape "
f"({self._num_envs}, 6, n), got {jac.shape}"
)
site_dof_indices = np.asarray(site.dof_vel_indices, dtype=np.int64).reshape(-1)
col_by_dof = {int(dof_index): col for col, dof_index in enumerate(site_dof_indices)}
requested = np.asarray(dof_indices, dtype=np.int64).reshape(-1)
cols: list[int] = []
for dof_index in requested:
key = int(dof_index)
if key not in col_by_dof:
raise ValueError(f"DoF index {key} is not present in site {sid} Jacobian")
cols.append(col_by_dof[key])
selected = jac[:, :, np.asarray(cols, dtype=np.intp)]
# Motrix returns angular rows first and linear rows second.
jacp = selected[:, 3:6, :]
jacr = selected[:, 0:3, :]
return jacp, jacr
def _resolve_single_dof_joint(self, name: str):
jid = self._model.get_joint_index(name)
if jid is None or jid < 0:
raise ValueError(f"Joint '{name}' not found in Motrix model")
joint = self._model.joints[int(jid)]
if int(getattr(joint, "num_dof_vel", 1)) != 1:
raise ValueError(f"Joint '{name}' is not a single-DoF joint")
return joint
def _joint_dof_local_index(self, name: str, model_index: int, *, pos: bool) -> int:
all_indices = self._joint_dof_pos_indices if pos else self._joint_dof_vel_indices
matches = np.flatnonzero(all_indices == int(model_index))
if matches.size != 1:
space = "qpos" if pos else "qvel"
raise ValueError(f"Joint '{name}' {space} index {model_index} is not in joint DoFs")
return int(matches[0])
[docs]
def get_geom_id(self, name: str) -> int:
geom_id = self._model.get_geom_index(name)
if geom_id is None or geom_id < 0:
raise ValueError(f"Geom '{name}' not found in Motrix model")
return int(geom_id)
[docs]
def get_geom_size(self, name: str) -> np.ndarray:
geom = _require_not_none(
self._model.get_geom(name),
f"Geom '{name}' not found in Motrix model",
)
return np.asarray(geom.size, dtype=np.float64).copy()
[docs]
def get_body_mass(self) -> np.ndarray:
return self._default_body_mass.copy()
[docs]
def get_body_ipos(self) -> np.ndarray:
return self._default_body_ipos.copy()
[docs]
def get_body_subtree_ids(self, root_body_id: int) -> np.ndarray:
root_id = int(root_body_id)
if root_id < 0 or root_id >= int(self._model.num_links):
raise ValueError(f"root_body_id out of range: {root_id}")
if root_id != int(self._body_link.index):
raise NotImplementedError(
"MotrixBackend only exposes the configured base articulation subtree"
)
subtree_ids = {root_id}
for link in self._model.links:
link_id = int(link.index)
joint_indices = getattr(link, "joint_indices", ())
if link_id != root_id and len(joint_indices) > 0:
subtree_ids.add(link_id)
return np.asarray(sorted(subtree_ids), dtype=np.int32)
[docs]
def get_geom_names(self) -> tuple[str, ...]:
return tuple(
str(getattr(self._geoms_by_id[geom_id], "name", "") or "")
for geom_id in range(int(self._model.num_geoms))
)
[docs]
def get_geom_body_ids(self) -> np.ndarray:
body_ids = np.zeros((int(self._model.num_geoms),), dtype=np.int32)
for geom_id in range(int(self._model.num_geoms)):
link = getattr(self._geoms_by_id[geom_id], "link", None)
if link is None:
body_ids[geom_id] = -1
else:
body_ids[geom_id] = int(link.index)
return body_ids
[docs]
def get_geom_friction(self) -> np.ndarray:
if not self._supports_geom_friction_override:
raise NotImplementedError("Motrix geom friction override is not available")
return self._default_geom_friction.copy()
[docs]
def get_gravity(self) -> np.ndarray:
return np.asarray(self._model.options.gravity, dtype=np.float64).copy()
[docs]
def get_joint_range(self) -> np.ndarray | None:
"""Return single-DoF joint limits in backend DOF order.
Motrix stores the model-wide limits as a ``(2, num_dof)`` table,
whereas the UniLab backend contract exposes the MuJoCo-shaped
``(num_dof, 2)`` table. This is materialized once by ``Entity`` and
never queried from a task hot path.
"""
raw_limits = np.asarray(self._model.joint_limits, dtype=self._np_dtype)
if raw_limits.ndim != 2 or raw_limits.shape != (2, self.num_dof_vel):
raise ValueError(
"Motrix joint limits must have shape (2, num_dof); "
f"received {raw_limits.shape} for {self.num_dof_vel} DOFs"
)
return np.array(raw_limits.T, copy=True)
# ------------------------------------------------------------------ #
# Simulation control #
# ------------------------------------------------------------------ #
[docs]
def step(self, ctrl: np.ndarray, nsteps: int = 1) -> dict | None:
if self._pre_step_control_fn is not None:
return self._step_with_pre_step_control(ctrl, nsteps)
t0 = time.perf_counter()
self._data.actuator_ctrls = np.ascontiguousarray(ctrl)
set_ctrl_ms = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
if nsteps == 1:
self._model.step(self._data)
else:
self._model.step_n(self._data, nsteps)
physics_ms = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._refresh_link_pose_cache()
self._invalidate_link_velocity_cache()
refresh_cache_ms = (time.perf_counter() - t0) * 1000.0
return {
"timing": {
"set_ctrl_ms": set_ctrl_ms,
"physics_ms": physics_ms,
"refresh_cache_ms": refresh_cache_ms,
}
}
def _step_with_pre_step_control(
self, ctrl: np.ndarray, nsteps: int
) -> dict[str, dict[str, float]]:
set_ctrl_ms = 0.0
physics_ms = 0.0
refresh_cache_ms = 0.0
for _ in range(nsteps):
t0 = time.perf_counter()
native_ctrl = self._apply_pre_step_control(ctrl)
self._data.actuator_ctrls = np.ascontiguousarray(native_ctrl)
set_ctrl_ms += (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._model.step(self._data)
physics_ms += (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._refresh_link_pose_cache()
self._invalidate_link_velocity_cache()
refresh_cache_ms += (time.perf_counter() - t0) * 1000.0
return {
"timing": {
"set_ctrl_ms": set_ctrl_ms,
"physics_ms": physics_ms,
"refresh_cache_ms": refresh_cache_ms,
}
}
[docs]
def set_state(
self,
env_indices: np.ndarray,
qpos: np.ndarray,
qvel: np.ndarray,
randomization: ResetRandomizationPayload | None = None,
) -> dict | None:
if randomization is not None:
unsupported = self.get_dr_capabilities().get_unsupported_reset_terms(
randomization.requested_terms()
)
if unsupported:
raise NotImplementedError(
f"Motrix reset randomization does not support terms: {sorted(unsupported)}"
)
timing: dict[str, float] = {
"set_state_mask_ms": 0.0,
"set_state_data_slice_ms": 0.0,
"set_state_data_reset_ms": 0.0,
"set_state_clear_forces_ms": 0.0,
"set_state_geom_overrides_ms": 0.0,
"set_state_reset_rand_ms": 0.0,
"set_state_set_dof_vel_ms": 0.0,
"set_state_set_dof_pos_ms": 0.0,
"set_state_actuator_ctrl_ms": 0.0,
"set_state_forward_kinematic_ms": 0.0,
"set_state_refresh_pose_cache_ms": 0.0,
"set_state_invalidate_velocity_ms": 0.0,
"set_state_qpos_convert_ms": 0.0,
"set_state_pool_reset_ms": 0.0,
"set_state_state_scatter_ms": 0.0,
"set_state_reset_upload_ms": 0.0,
"set_state_reset_forward_ms": 0.0,
"set_state_host_cache_refresh_ms": 0.0,
"set_state_internal_gap_ms": 0.0,
}
outer_t0 = time.perf_counter()
# Pre-convert env_indices once; every downstream helper reuses this.
env_ids_intp = np.asarray(env_indices, dtype=np.intp)
t0 = time.perf_counter()
# Reuse the scratch qpos buffer when its shape matches; the reset path
# feeds a fixed-shape (num_envs, qpos_dim) array so this holds after
# the first call.
scratch = self._set_state_qpos_motrix_scratch
if scratch is None or scratch.shape != qpos.shape or scratch.dtype != qpos.dtype:
scratch = np.empty_like(qpos)
self._set_state_qpos_motrix_scratch = scratch
qpos_motrix = self._mujoco_qpos_to_motrix_into(scratch, qpos)
timing["set_state_qpos_convert_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
# Reuse the persistent bool-mask scratch; only the touched entries need
# rewriting each call.
mask = self._set_state_mask_scratch
if mask.shape[0] != self._num_envs:
mask = np.zeros(self._num_envs, dtype=bool)
self._set_state_mask_scratch = mask
mask.fill(False)
mask[env_ids_intp] = True
timing["set_state_mask_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
data_slice = self._data[mask]
timing["set_state_data_slice_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
data_slice.reset(self._model)
timing["set_state_data_reset_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._clear_applied_body_forces(env_indices, env_ids_intp=env_ids_intp)
timing["set_state_clear_forces_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._apply_init_geom_size_overrides(data_slice, env_indices, env_ids_intp=env_ids_intp)
timing["set_state_geom_overrides_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._apply_reset_randomization(
data_slice, env_indices, randomization, env_ids_intp=env_ids_intp
)
timing["set_state_reset_rand_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
data_slice.set_dof_vel(qvel)
timing["set_state_set_dof_vel_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
data_slice.set_dof_pos(qpos_motrix, self._model)
timing["set_state_set_dof_pos_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
if self._supports_position_actuator_gains and len(self._joint_dof_pos_indices) == int(
self.num_actuators
):
# Fully-actuated model: hold every joint at its reset position (unchanged).
if self._joint_dof_pos_slice is not None:
ctrl = qpos_motrix[:, self._joint_dof_pos_slice]
else:
ctrl = qpos_motrix[:, self._joint_dof_pos_indices]
elif self._actuator_joint_pos_indices is not None:
# Under-actuated / parallel model: hold only the actuated joints.
if self._actuator_joint_pos_slice is not None:
ctrl = qpos_motrix[:, self._actuator_joint_pos_slice]
else:
ctrl = qpos_motrix[:, self._actuator_joint_pos_indices]
else:
ctrl = np.zeros((len(env_indices), self.num_actuators), dtype=self._np_dtype)
# Only pay the copy when the underlying slice is non-contiguous; view
# slices of a contiguous scratch buffer already satisfy the motrixsim
# contiguous requirement.
if not ctrl.flags.c_contiguous:
ctrl = np.ascontiguousarray(ctrl)
data_slice.actuator_ctrls = ctrl
timing["set_state_actuator_ctrl_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._model.forward_kinematic(data_slice)
timing["set_state_forward_kinematic_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._refresh_link_pose_cache(env_indices, data_slice=data_slice, env_ids_intp=env_ids_intp)
timing["set_state_refresh_pose_cache_ms"] = (time.perf_counter() - t0) * 1000.0
t0 = time.perf_counter()
self._invalidate_link_velocity_cache()
timing["set_state_invalidate_velocity_ms"] = (time.perf_counter() - t0) * 1000.0
outer_total_ms = (time.perf_counter() - outer_t0) * 1000.0
measured_ms = (
timing["set_state_qpos_convert_ms"]
+ timing["set_state_mask_ms"]
+ timing["set_state_data_slice_ms"]
+ timing["set_state_data_reset_ms"]
+ timing["set_state_clear_forces_ms"]
+ timing["set_state_geom_overrides_ms"]
+ timing["set_state_reset_rand_ms"]
+ timing["set_state_set_dof_vel_ms"]
+ timing["set_state_set_dof_pos_ms"]
+ timing["set_state_actuator_ctrl_ms"]
+ timing["set_state_forward_kinematic_ms"]
+ timing["set_state_refresh_pose_cache_ms"]
+ timing["set_state_invalidate_velocity_ms"]
)
timing["set_state_internal_gap_ms"] = outer_total_ms - measured_ms
return {"timing": timing}
[docs]
def get_dr_capabilities(self) -> DomainRandomizationCapabilities:
supported_reset_terms = {
RESET_TERM_BASE_MASS,
RESET_TERM_BASE_COM,
RESET_TERM_BODY_MASS,
RESET_TERM_BODY_IPOS,
}
if getattr(self, "_supports_position_actuator_gains", False):
supported_reset_terms.update({RESET_TERM_KP, RESET_TERM_KD})
if getattr(self, "_supports_geom_friction_override", False):
supported_reset_terms.add(RESET_TERM_GEOM_FRICTION)
if getattr(self, "_supports_gravity_override", False):
supported_reset_terms.add(RESET_TERM_GRAVITY)
return DomainRandomizationCapabilities(
supported_reset_terms=frozenset(supported_reset_terms),
supports_interval_push=True,
supports_interval_body_velocity_delta=False,
supports_interval_body_force=getattr(self, "_supports_external_force", False),
supported_interval_terms=frozenset(
{INTERVAL_TERM_PUSH}
| (
{INTERVAL_TERM_BODY_FORCE}
if getattr(self, "_supports_external_force", False)
else set()
)
),
)
[docs]
def apply_init_randomization(self, plan: InitRandomizationPlan) -> None:
if plan.is_empty():
return
model_assignments = np.asarray(plan.model_assignments, dtype=np.int32)
if model_assignments.shape != (self._num_envs,):
raise ValueError(
f"model_assignments must have shape ({self._num_envs},), "
f"got {model_assignments.shape}"
)
if np.any(model_assignments < 0) or np.any(model_assignments >= len(plan.model_variants)):
raise ValueError(
"model_assignments must refer to entries in InitRandomizationPlan.model_variants"
)
geom_size_overrides: dict[int, np.ndarray] = {}
for variant_id, variant in enumerate(plan.model_variants):
env_indices = np.flatnonzero(model_assignments == variant_id)
if env_indices.size == 0:
continue
for override in variant.geom_size_overrides:
geom_id = self.get_geom_id(override.geom_name)
geom = _require_not_none(
self._model.get_geom(geom_id),
f"Geom '{override.geom_name}' not found in Motrix model",
)
override_shape = np.asarray(geom.get_size_override(self._data)).shape
if len(override_shape) != 2:
raise ValueError(
f"Motrix geom '{override.geom_name}' size override must be rank-2, "
f"got shape {override_shape}"
)
width = int(override_shape[1])
size = np.asarray(override.size, dtype=np.float64).reshape(-1)
if size.size < width:
raise ValueError(
f"GeomSizeOverride for '{override.geom_name}' has {size.size} values, "
f"but Motrix expects at least {width}"
)
values = geom_size_overrides.setdefault(
geom_id,
np.asarray(geom.get_size_override(self._data), dtype=np.float64).copy(),
)
values[env_indices, :] = size[:width]
self._init_geom_size_overrides = geom_size_overrides
self._apply_init_geom_size_overrides(self._data, np.arange(self._num_envs, dtype=np.int32))
_interval_term_handler_cache: dict[str, Callable[[IntervalTermOp], None]] | None = None
def _interval_term_handlers(self) -> dict[str, Callable[[IntervalTermOp], None]]:
# Built lazily once. Torque, angular-velocity and linear-velocity
# terms intentionally have no handler and fail closed in the base
# dispatch; body force stays gated inside ``apply_body_force`` on the
# runtime external-force API probe.
if self._interval_term_handler_cache is None:
self._interval_term_handler_cache = {
INTERVAL_TERM_PUSH: lambda op: self.push_robots(op.payload),
INTERVAL_TERM_BODY_FORCE: lambda op: self.apply_body_force(
op.body_ids, op.payload
),
}
return self._interval_term_handler_cache
[docs]
def get_play_capabilities(self) -> BackendPlayCapabilities:
return BackendPlayCapabilities(
supports_native_interactive_renderer=True,
supports_native_video_capture=True,
)
[docs]
def resolve_play_render_plan(
self,
*,
play_render_mode: str | None,
play_steps: int | None,
output_video: str | os.PathLike[str] | None,
) -> BackendPlayRenderPlan:
mode = normalize_play_render_mode(play_render_mode)
effective_mode = "interactive" if mode == "auto" else mode
if effective_mode == "none":
return BackendPlayRenderPlan(
mode=effective_mode,
headless=True,
record_video=False,
num_steps=None,
output_video=None,
)
if effective_mode == "interactive":
return BackendPlayRenderPlan(
mode=effective_mode,
headless=False,
record_video=False,
num_steps=None,
output_video=None,
)
assert effective_mode == "record"
if play_steps is None:
raise ValueError("Motrix record playback requires a finite training.play_steps value.")
if output_video is None:
raise ValueError("Motrix record playback requires an output video path.")
return BackendPlayRenderPlan(
mode=effective_mode,
headless=True,
record_video=True,
num_steps=int(play_steps),
output_video=output_video,
)
[docs]
def run_playback(
self,
*,
env: Any,
initialize,
step,
num_steps: int | None,
output_video: str | os.PathLike[str] | None = None,
render_spacing: float | None = None,
render_offset_mode: str | None = None,
headless: bool | None = None,
record_video: bool | None = None,
frame_state_getter=None,
camera_kwargs: CameraCfg | Mapping[str, Any] | None = None,
debug_overlay_getter=None,
on_frame=None,
) -> str | None:
del frame_state_getter
if debug_overlay_getter is not None:
raise unsupported_debug_overlay_error(self.__class__.__name__)
if on_frame is not None:
raise NotImplementedError(
f"{self.__class__.__name__} renders through a native renderer and "
"does not support on_frame callbacks"
)
camera = CameraCfg.from_kwargs(camera_kwargs)
should_record_video = (
bool(record_video) if record_video is not None else output_video is not None
)
should_run_headless = bool(headless) if headless is not None else should_record_video
try:
return run_motrix_playback(
backend=self,
env=env,
initialize=initialize,
step=step,
num_steps=num_steps,
output_video=output_video,
render_spacing=render_spacing,
render_offset_mode=render_offset_mode,
headless=should_run_headless,
record_video=should_record_video,
camera_kwargs=camera,
)
except RenderClosedError:
if not should_run_headless and not should_record_video:
logger.info("Render window closed.")
return None
raise
# ------------------------------------------------------------------ #
# Base kinematics #
# ------------------------------------------------------------------ #
[docs]
def get_base_pos(self) -> np.ndarray:
if self._body_floatingbase is not None:
return self._body_floatingbase.get_translation(self._data) # type: ignore[no-any-return]
return self._body_link.get_pose(self._data)[:, :3] # type: ignore[no-any-return]
[docs]
def get_base_quat(self) -> np.ndarray:
if self._body_floatingbase is not None:
quat = self._body_floatingbase.get_rotation(self._data)
else:
quat = self._body_link.get_rotation(self._data)
return self._xyzw_to_wxyz(quat)
[docs]
def get_base_lin_vel(self) -> np.ndarray:
if self._body_floatingbase is not None:
return self._body_floatingbase.get_global_linear_velocity(self._data) # type: ignore[no-any-return]
return self._body_link.get_linear_velocity(self._data) # type: ignore[no-any-return]
[docs]
def get_base_ang_vel(self) -> np.ndarray:
if self._body_floatingbase is not None:
return self._body_floatingbase.get_global_angular_velocity(self._data) # type: ignore[no-any-return]
return self._body_link.get_angular_velocity(self._data) # type: ignore[no-any-return]
# ------------------------------------------------------------------ #
# DOF state #
# ------------------------------------------------------------------ #
[docs]
def get_dof_pos(self) -> np.ndarray:
indices = (
self._actuator_joint_pos_indices
if self._actuator_joint_pos_indices is not None
else self._joint_dof_pos_indices
)
result = self._data.dof_pos[..., indices] # type: ignore[no-any-return]
return result
[docs]
def get_dof_vel(self) -> np.ndarray:
indices = (
self._actuator_joint_vel_indices
if self._actuator_joint_vel_indices is not None
else self._joint_dof_vel_indices
)
return self._data.dof_vel[..., indices] # type: ignore[no-any-return]
# ------------------------------------------------------------------ #
# Body kinematics — world frame #
# ------------------------------------------------------------------ #
def _as_body_ids(self, body_ids: np.ndarray) -> np.ndarray:
return np.asarray(body_ids, dtype=np.int32)
[docs]
def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray:
return self._get_link_poses_w(body_ids)[:, :, :3]
[docs]
def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray:
return self._xyzw_to_wxyz(self._get_link_poses_w(body_ids)[:, :, 3:])
[docs]
def get_body_pose_w_rows(
self, env_ids: np.ndarray, body_ids: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
rows = np.asarray(env_ids, dtype=np.intp)
poses_w = self._link_poses[rows[:, None], self._as_body_ids(body_ids), :]
return poses_w[:, :, :3], self._xyzw_to_wxyz(poses_w[:, :, 3:])
[docs]
def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
poses = self._get_link_poses_w(body_ids)
return poses[:, :, :3], self._xyzw_to_wxyz(poses[:, :, 3:])
[docs]
def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray:
return self._get_link_lin_vel_w(body_ids)
[docs]
def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray:
return self._get_link_ang_vel_w(body_ids)
[docs]
def get_body_state_w(
self, body_ids: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
poses_w = self._get_link_poses_w(body_ids)
lin_vel_w, ang_vel_w = self.get_body_vel_w(body_ids)
return (
poses_w[:, :, :3],
self._xyzw_to_wxyz(poses_w[:, :, 3:]),
lin_vel_w,
ang_vel_w,
)
[docs]
def copy_body_state_w(
self,
body_ids: np.ndarray,
out_pos: np.ndarray,
out_quat: np.ndarray,
out_lin_vel: np.ndarray,
out_ang_vel: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
ids = self._as_body_ids(body_ids)
poses_w = self._get_link_poses_w(ids)
out_pos[..., 0] = poses_w[..., 0]
out_pos[..., 1] = poses_w[..., 1]
out_pos[..., 2] = poses_w[..., 2]
out_quat[..., 0] = poses_w[..., 6]
out_quat[..., 1] = poses_w[..., 3]
out_quat[..., 2] = poses_w[..., 4]
out_quat[..., 3] = poses_w[..., 5]
link_velocity_cache = self._ensure_link_velocity_cache()
if self._link_velocity_cache is None or self._link_velocity_cache.shape != (
self._num_envs,
len(ids),
6,
):
self._link_velocity_cache = np.empty(
(self._num_envs, len(ids), 6), dtype=self._np_dtype
)
np.take(link_velocity_cache, ids, axis=1, out=self._link_velocity_cache)
out_lin_vel[..., 0] = self._link_velocity_cache[..., 0]
out_lin_vel[..., 1] = self._link_velocity_cache[..., 1]
out_lin_vel[..., 2] = self._link_velocity_cache[..., 2]
out_ang_vel[..., 0] = self._link_velocity_cache[..., 3]
out_ang_vel[..., 1] = self._link_velocity_cache[..., 4]
out_ang_vel[..., 2] = self._link_velocity_cache[..., 5]
return out_pos, out_quat, out_lin_vel, out_ang_vel
[docs]
def get_body_vel_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
ids = self._as_body_ids(body_ids)
velocities = np.ascontiguousarray(self._ensure_link_velocity_cache()[:, ids, :])
return velocities[:, :, :3], velocities[:, :, 3:]
[docs]
def get_body_lin_vel_w_rows(self, env_ids: np.ndarray, body_ids: np.ndarray) -> np.ndarray:
rows = np.asarray(env_ids, dtype=np.intp)
return self._ensure_link_velocity_cache()[rows[:, None], self._as_body_ids(body_ids), :3] # type: ignore[no-any-return]
[docs]
def get_body_ang_vel_w_rows(self, env_ids: np.ndarray, body_ids: np.ndarray) -> np.ndarray:
rows = np.asarray(env_ids, dtype=np.intp)
return self._ensure_link_velocity_cache()[rows[:, None], self._as_body_ids(body_ids), 3:] # type: ignore[no-any-return]
# ------------------------------------------------------------------ #
# Body kinematics — baselink frame #
# ------------------------------------------------------------------ #
[docs]
def get_body_pos_b(self, body_ids: np.ndarray) -> np.ndarray:
return self._get_body_sensor_values(body_ids, "track_pos_b")
[docs]
def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray:
# MotrixSim framequat sensors output xyzw; convert to the wxyz contract.
return self._xyzw_to_wxyz(self._get_body_sensor_values(body_ids, "track_quat_b"))
[docs]
def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray:
# Analytical per the SimBackend contract: world-frame velocity rotated
# into each body's own frame. MotrixSim frame sensors report motion
# relative to the baselink and degenerate to zero for the root body.
ids = self._as_body_ids(body_ids)
return np_quat_apply_inverse_batched(
self.get_body_quat_w(ids), self._get_link_lin_vel_w(ids)
)
[docs]
def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray:
ids = self._as_body_ids(body_ids)
return np_quat_apply_inverse_batched(
self.get_body_quat_w(ids), self._get_link_ang_vel_w(ids)
)
# ------------------------------------------------------------------ #
# Sensors #
# ------------------------------------------------------------------ #
def _validate_sensor_names(self, names: Sequence[str]) -> tuple[str, ...]:
sensor_names = tuple(names)
missing = tuple(name for name in sensor_names if name not in self._sensor_names)
if missing:
available = ", ".join(sorted(self._sensor_names))
raise KeyError(f"Unknown Motrix sensor(s) {missing}; available sensors: {available}")
return sensor_names
[docs]
def get_sensor_data(self, name: str) -> np.ndarray:
self._validate_sensor_names((name,))
return self._model.get_sensor_value(name, self._data) # type: ignore[no-any-return]
[docs]
def get_sensor_data_rows(self, name: str, env_ids: np.ndarray) -> np.ndarray:
self._validate_sensor_names((name,))
rows = np.asarray(env_ids, dtype=np.intp)
mask = np.zeros(self._num_envs, dtype=bool)
mask[rows] = True
selected_rows = np.flatnonzero(mask)
selected_values = self._model.get_sensor_value(name, self._data[mask])
return selected_values[np.searchsorted(selected_rows, rows)] # type: ignore[no-any-return]
[docs]
def get_sensor_data_batch(self, names: Sequence[str]) -> np.ndarray:
sensor_names = self._validate_sensor_names(names)
if not sensor_names:
return np.empty((self._num_envs, 0), dtype=self._np_dtype)
values = self._model.get_sensor_values(sensor_names, self._data)
return np.asarray(values, dtype=self._np_dtype)
def _bind_sensor_data_reader(self, names: tuple[str, ...]) -> Callable[[], np.ndarray]:
"""Retain Motrix's opaque native reader after cold-path name validation.
MotrixSim exposes named sensor access but no public numeric sensor ID;
the bound callable is therefore the narrowest backend-owned reader.
It does not inspect XML or model metadata on the manager hot path.
"""
native_reader = self._model.get_sensor_values
def read() -> np.ndarray:
return np.asarray(native_reader(names, self._data), dtype=self._np_dtype)
return read
# ------------------------------------------------------------------ #
# MotrixSim-specific #
# ------------------------------------------------------------------ #
def _get_body_names(self, body_ids: np.ndarray) -> list[str]:
return [self._body_id_to_name[int(bid)] for bid in self._as_body_ids(body_ids)]
def _get_link_poses_w(self, body_ids: np.ndarray) -> np.ndarray:
ids = self._as_body_ids(body_ids)
return np.ascontiguousarray(self._link_poses[:, ids, :])
def _get_link_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray:
ids = self._as_body_ids(body_ids)
return np.ascontiguousarray(self._ensure_link_velocity_cache()[:, ids, :3])
def _get_link_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray:
ids = self._as_body_ids(body_ids)
return np.ascontiguousarray(self._ensure_link_velocity_cache()[:, ids, 3:])
def _get_body_sensor_values(self, body_ids: np.ndarray, prefix: str) -> np.ndarray:
return np.stack(
[
self._model.get_sensor_value(f"{prefix}_{name}", self._data)
for name in self._get_body_names(body_ids)
],
axis=1,
)
def _xyzw_to_wxyz(self, q: np.ndarray) -> np.ndarray:
"""motrix xyzw → wxyz"""
return q[..., [3, 0, 1, 2]]
def _mujoco_qpos_to_motrix(self, qpos: np.ndarray) -> np.ndarray:
"""Convert every MuJoCo freejoint quaternion slice from wxyz to xyzw."""
qpos_motrix = np.array(qpos, copy=True)
for quat_indices in self._floating_base_quat_indices:
qpos_motrix[..., quat_indices] = qpos[..., quat_indices[[1, 2, 3, 0]]]
return qpos_motrix
def _mujoco_qpos_to_motrix_into(self, dst: np.ndarray, qpos: np.ndarray) -> np.ndarray:
"""Hot-path variant that writes into ``dst`` in place.
Same conversion as :meth:`_mujoco_qpos_to_motrix` but avoids the per-call
``np.array(qpos, copy=True)`` allocation. Returns ``dst`` so callers can
chain: ``qpos_motrix = self._mujoco_qpos_to_motrix_into(scratch, qpos)``.
``dst`` must have the same shape as ``qpos``; caller is responsible for
sizing the scratch buffer.
"""
np.copyto(dst, qpos, casting="same_kind")
for quat_indices in self._floating_base_quat_indices:
dst[..., quat_indices] = qpos[..., quat_indices[[1, 2, 3, 0]]]
return dst
def _motrix_qpos_to_mujoco(self, qpos: np.ndarray) -> np.ndarray:
"""Convert every Motrix freejoint quaternion slice from xyzw to wxyz."""
qpos_mujoco = np.array(qpos, copy=True)
for quat_indices in self._floating_base_quat_indices:
qpos_mujoco[..., quat_indices] = qpos[..., quat_indices[[3, 0, 1, 2]]]
return qpos_mujoco
def _refresh_link_pose_cache(
self,
env_indices: np.ndarray | None = None,
data_slice: Any | None = None,
env_ids_intp: np.ndarray | None = None,
) -> None:
if env_indices is None:
self._link_poses = self._model.get_link_poses(self._data)
else:
if data_slice is None:
mask = np.zeros(self._num_envs, dtype=bool)
mask[env_indices] = True
data_slice = self._data[mask]
# Indexing with intp is a hair faster than the arbitrary-dtype
# ndarray path; use the pre-converted array when set_state hands
# it down.
idx = env_indices if env_ids_intp is None else env_ids_intp
self._link_poses[idx] = self._model.get_link_poses(data_slice)
def _refresh_link_velocity_cache(self, env_indices: np.ndarray | None = None) -> None:
if env_indices is None:
self._link_velocities = self._model.get_link_velocities(self._data)
else:
mask = np.zeros(self._num_envs, dtype=bool)
mask[env_indices] = True
if self._link_velocities is None:
self._link_velocities = self._model.get_link_velocities(self._data)
self._link_velocity_cache_valid = True
return
self._link_velocities[env_indices] = self._model.get_link_velocities(self._data[mask])
self._link_velocity_cache_valid = True
def _invalidate_link_velocity_cache(self) -> None:
self._link_velocity_cache_valid = False
def _ensure_link_velocity_cache(self) -> np.ndarray:
if self._link_velocities is None or not self._link_velocity_cache_valid:
self._refresh_link_velocity_cache()
assert self._link_velocities is not None
return self._link_velocities
def _coerce_reset_field(
self,
value: np.ndarray,
*,
name: str,
num_reset: int,
shaped_tail: tuple[int, ...],
) -> np.ndarray:
arr = np.asarray(value, dtype=np.float32)
shaped = (num_reset, *shaped_tail)
flat_shape = (num_reset, int(np.prod(shaped_tail)))
if arr.shape == shaped:
return arr.copy()
if arr.shape == flat_shape:
return arr.reshape(shaped).copy()
raise ValueError(f"{name} must have shape {shaped} or {flat_shape}, got {arr.shape}")
def _apply_init_geom_size_overrides(
self,
data_slice,
env_indices: np.ndarray,
env_ids_intp: np.ndarray | None = None,
) -> None:
if not self._init_geom_size_overrides:
return
env_ids = (
env_ids_intp if env_ids_intp is not None else np.asarray(env_indices, dtype=np.intp)
)
for geom_id, values in self._init_geom_size_overrides.items():
geom = _require_not_none(
self._model.get_geom(int(geom_id)),
f"Geom id {geom_id} not found in Motrix model",
)
geom.set_size_override(
data_slice,
np.ascontiguousarray(np.asarray(values[env_ids], dtype=np.float32)),
)
def _set_link_mass_overrides(self, data_slice, body_mass: np.ndarray) -> None:
for link_id, link in self._links_by_id.items():
link.set_mass_override(
data_slice,
np.ascontiguousarray(np.asarray(body_mass[:, link_id], dtype=np.float32)),
)
def _set_link_ipos_overrides(self, data_slice, body_ipos: np.ndarray) -> None:
for link_id, link in self._links_by_id.items():
link.set_center_of_mass_override(
data_slice,
np.ascontiguousarray(np.asarray(body_ipos[:, link_id, :], dtype=np.float32)),
)
def _set_geom_friction_overrides(self, data_slice, geom_friction: np.ndarray) -> None:
if not self._supports_geom_friction_override:
raise NotImplementedError("Motrix geom friction override is not available")
override_ids = getattr(self, "_geom_friction_override_ids", tuple(self._geoms_by_id))
unsupported_ids = sorted(set(self._geoms_by_id) - set(override_ids))
if unsupported_ids:
unsupported_values = geom_friction[:, unsupported_ids, :]
default_values = self._default_geom_friction[None, unsupported_ids, :]
if not np.allclose(unsupported_values, default_values):
raise ValueError(
"Motrix geom friction override only supports collision geoms; "
f"non-collision geom ids were modified: {unsupported_ids}"
)
for geom_id in override_ids:
geom = self._geoms_by_id[int(geom_id)]
geom.set_friction_override(
data_slice,
np.ascontiguousarray(np.asarray(geom_friction[:, geom_id, :], dtype=np.float32)),
)
def _clear_applied_body_forces(
self,
env_indices: np.ndarray,
env_ids_intp: np.ndarray | None = None,
) -> None:
if not self._applied_body_forces:
return
env_ids = (
env_ids_intp if env_ids_intp is not None else np.asarray(env_indices, dtype=np.intp)
)
for applied_force in self._applied_body_forces.values():
applied_force[env_ids, :] = 0.0
[docs]
def push_robots(self, force_range):
ex_force = np.random.rand(self.num_envs, 3) * 2 - 1 # [x_force, y_force, z_force]
ex_force[:, 0] *= force_range[0]
ex_force[:, 1] *= force_range[1]
ex_force[:, 2] *= force_range[2]
self._push_body_link.add_external_force(self._data, ex_force, local=True)
[docs]
def apply_body_force(
self,
body_ids: np.ndarray,
force: np.ndarray,
torque: np.ndarray | None = None,
) -> None:
"""Apply absolute world-frame external forces through Motrix Link API."""
if torque is not None:
raise NotImplementedError(
f"{self.__class__.__name__} does not support interval body torque perturbation"
)
if not getattr(self, "_supports_external_force", False):
raise NotImplementedError("Motrix link external-force API is not available")
body_ids_np = np.asarray(body_ids, dtype=np.int32).reshape(-1)
force_np = np.asarray(force, dtype=np.float64)
expected_shape = (self._num_envs, body_ids_np.size, 3)
if force_np.shape != expected_shape:
raise ValueError(f"body force must have shape {expected_shape}, got {force_np.shape}")
for body_offset, body_id in enumerate(body_ids_np):
link_id = int(body_id)
link = self._links_by_id.get(link_id)
if link is None:
raise ValueError(f"Body id {link_id} not found in Motrix model")
target_force = np.asarray(force_np[:, body_offset, :], dtype=np.float64)
applied_force = self._applied_body_forces.setdefault(
link_id,
np.zeros((self._num_envs, 3), dtype=np.float64),
)
delta_force = target_force - applied_force
if np.any(delta_force):
link.add_external_force(
self._data,
np.ascontiguousarray(delta_force.astype(np.float32)),
local=False,
)
applied_force[:] = target_force
[docs]
def create_hfield_scanner(
self,
*,
hfield_geom_id: int,
offsets: np.ndarray,
frame_body_id: int,
alignment: str = "yaw",
output: str = "height",
) -> BackendHeightScanner:
offsets_np = np.ascontiguousarray(np.asarray(offsets, dtype=np.float32))
if offsets_np.ndim != 2 or offsets_np.shape[1] != 2:
raise ValueError(f"offsets must have shape (num_points, 2), got {offsets_np.shape}")
if alignment != "yaw":
raise ValueError(f"MotrixBackend only supports alignment='yaw', got {alignment!r}")
if output not in {"height", "clearance"}:
raise ValueError(f"Unsupported hfield sampling output: {output!r}")
geom_id = int(hfield_geom_id)
if geom_id < 0 or geom_id >= int(self._model.num_geoms):
raise ValueError(f"hfield_geom_id out of range: {geom_id}")
body_id = int(frame_body_id)
if body_id < 0 or body_id >= int(self._model.num_links):
raise ValueError(f"frame_body_id out of range: {body_id}")
terrain = self._model.get_geom(geom_id)
if terrain is None:
raise ValueError(f"Geom id {geom_id} not found in Motrix model")
if not isinstance(terrain, mtx.GeomHField):
raise ValueError(f"Geom id {geom_id} is not backed by a Motrix hfield")
frame = self._link_cache[body_id]
scanner = mtx.TerrainScanner(
terrain,
frame,
offsets_np,
alignment=alignment,
output=output,
)
return _MotrixTerrainScanner(
scanner=scanner,
data=self._data,
out=np.empty((self._num_envs, offsets_np.shape[0]), dtype=self._np_dtype),
)
def _update_tracking_camera_view(self) -> None:
if (
self._render_app is None
or self._render_tracking_camera is None
or self._render_offsets_np is None
):
return
lookat = tracking_camera_lookat(
self.get_base_pos(),
self._render_tracking_camera,
self._render_offsets_np,
)
self._render_app.system_camera.set_view(
lookat,
self._render_tracking_camera.distance,
self._render_tracking_camera.elevation,
self._render_tracking_camera.azimuth,
)
def _assert_render_context_available(self, *, headless: bool, capture: bool) -> None:
if self._render_app is None:
return
if self._render_headless != headless:
raise RuntimeError(
"Motrix renderer is already initialized with "
f"headless={self._render_headless!r}; cannot reuse it with headless={headless!r}"
)
if capture and not self._render_capture_enabled:
raise RuntimeError(
"Motrix renderer is already initialized without video capture; "
"cannot enable capture on the existing renderer"
)
return
[docs]
def init_renderer(
self,
spacing: float = 1.0,
*,
offset_mode: str = "grid",
headless: bool = False,
capture: bool = False,
width: int = 1280,
height: int = 720,
camera_kwargs: CameraCfg | Mapping[str, Any] | None = None,
) -> None:
"""Initialize a Motrix renderer, optionally enabling system-camera capture."""
headless = bool(headless)
capture = bool(capture)
self._assert_render_context_available(headless=headless, capture=capture)
if self._render_app is not None:
return
camera = CameraCfg.from_kwargs(camera_kwargs)
settings = RenderSettings.performance()
settings.enable_shadow = True
offsets = render_offsets(
self._num_envs,
float(spacing),
offset_mode=str(offset_mode),
)
offsets_np = np.asarray(offsets, dtype=np.float64)
self._render_offsets_np = offsets_np
use_configured_camera = capture or camera_kwargs is not None
if use_configured_camera:
base_positions = self.get_base_pos() if camera.cam_tracking else None
camera_view = resolve_system_camera_view(
self._num_envs,
base_positions,
offsets,
camera,
)
tracking_camera = camera_view.tracking
else:
tracking_camera = None
if capture:
self._model.cameras.set_system_render_target("image", int(width), int(height))
render_app = RenderApp(headless=headless)
try:
render_app.launch(
self._model,
batch=self._num_envs,
render_offset=offsets,
render_settings=settings,
)
except _MotrixRenderClosedError as e:
# Normalize the motrixsim-private window-closed error to the
# interface-level signal declared on SimBackend.
raise RenderClosedError(str(e)) from e
if use_configured_camera:
render_app.system_camera.set_view(
camera_view.lookat,
camera_view.distance,
camera_view.elevation,
camera_view.azimuth,
)
if not capture:
render_app.set_main_camera(None)
self._render_app = render_app
self._render_headless = headless
self._render_capture_enabled = capture
self._render_tracking_camera = tracking_camera
[docs]
def render(self):
"""Render current state (interactive visualization)"""
if self._render_app is None:
self.init_renderer()
self._assert_render_context_available(headless=False, capture=False)
assert self._render_app is not None
self._update_tracking_camera_view()
try:
self._render_app.sync(data=self._data)
except _MotrixRenderClosedError as e:
# Normalize the motrixsim-private window-closed error to the
# interface-level signal declared on SimBackend.
raise RenderClosedError(str(e)) from e
[docs]
def capture_video_frame(self) -> np.ndarray:
"""Capture one RGB frame from Motrix's system camera."""
if self._render_app is None:
self.init_renderer(headless=True, capture=True)
if not self._render_capture_enabled:
raise RuntimeError("Motrix renderer is not initialized for video capture")
assert self._render_app is not None
self._update_tracking_camera_view()
try:
task = self._render_app.system_camera.capture()
self._render_app.sync(data=self._data, wait=True)
image = task.take_image()
except _MotrixRenderClosedError as e:
# Normalize the motrixsim-private window-closed error to the
# interface-level signal declared on SimBackend.
raise RenderClosedError(str(e)) from e
if image is None:
raise RuntimeError("Motrix system camera capture did not return an image")
pixels = np.asarray(image.pixels)
if pixels.ndim != 3:
raise RuntimeError(
f"Motrix system camera capture must return an HWC image, got shape {pixels.shape}"
)
if pixels.shape[-1] == 4:
pixels = pixels[..., :3]
if pixels.shape[-1] != 3:
raise RuntimeError(
"Motrix system camera capture must return RGB/RGBA pixels, "
f"got shape {pixels.shape}"
)
return np.ascontiguousarray(pixels, dtype=np.uint8)
def _apply_reset_randomization(
self,
data_slice,
env_indices: np.ndarray,
randomization: ResetRandomizationPayload | None,
env_ids_intp: np.ndarray | None = None,
) -> None:
if randomization is None or randomization.is_empty():
return
unsupported = (
randomization.requested_terms() - self.get_dr_capabilities().supported_reset_terms
)
if unsupported:
terms = ", ".join(sorted(unsupported))
raise NotImplementedError(
f"{self.backend_type} backend does not support reset randomization terms: {terms}"
)
env_ids = (
env_ids_intp if env_ids_intp is not None else np.asarray(env_indices, dtype=np.intp)
)
num_reset = len(env_ids)
body_mass = None
if randomization.body_mass is not None:
body_mass = self._coerce_reset_field(
randomization.body_mass,
name="body_mass",
num_reset=num_reset,
shaped_tail=(int(self._model.num_links),),
)
if randomization.base_mass_delta is not None:
if body_mass is None:
body_mass = np.broadcast_to(
self._default_body_mass,
(num_reset, int(self._model.num_links)),
).copy()
body_mass[:, int(self._body_link.index)] += np.asarray(
randomization.base_mass_delta, dtype=np.float32
)
if body_mass is not None:
self._set_link_mass_overrides(data_slice, body_mass)
body_ipos = None
if randomization.body_ipos is not None:
body_ipos = self._coerce_reset_field(
randomization.body_ipos,
name="body_ipos",
num_reset=num_reset,
shaped_tail=(int(self._model.num_links), 3),
)
if randomization.base_com_offset is not None:
if body_ipos is None:
body_ipos = np.broadcast_to(
self._default_body_ipos,
(num_reset, int(self._model.num_links), 3),
).copy()
body_ipos[:, int(self._body_link.index), :] += np.asarray(
randomization.base_com_offset, dtype=np.float32
)
if body_ipos is not None:
self._set_link_ipos_overrides(data_slice, body_ipos)
if randomization.geom_friction is not None:
geom_friction = self._coerce_reset_field(
randomization.geom_friction,
name="geom_friction",
num_reset=num_reset,
shaped_tail=(int(self._model.num_geoms), 3),
)
self._set_geom_friction_overrides(data_slice, geom_friction)
if randomization.gravity is not None:
gravity = self._coerce_reset_field(
randomization.gravity,
name="gravity",
num_reset=num_reset,
shaped_tail=(3,),
)
self._set_gravity_override(data_slice, gravity)
if randomization.kp is not None:
kp = np.asarray(randomization.kp, dtype=np.float32)
expected_shape = (num_reset, self.num_actuators)
if kp.shape != expected_shape:
raise ValueError(f"kp must have shape {expected_shape}, got {kp.shape}")
self._set_position_actuator_kp_override(data_slice, kp)
if randomization.kd is not None:
kd = np.asarray(randomization.kd, dtype=np.float32)
expected_shape = (num_reset, self.num_actuators)
if kd.shape != expected_shape:
raise ValueError(f"kd must have shape {expected_shape}, got {kd.shape}")
self._set_position_actuator_kd_override(data_slice, kd)
def _set_gravity_override(self, data_slice, gravity: np.ndarray) -> None:
if not getattr(self, "_supports_gravity_override", False):
raise NotImplementedError("Motrix gravity override is not available")
self._model.set_gravity_override(
data_slice,
np.ascontiguousarray(np.asarray(gravity, dtype=np.float32)),
)
def _set_position_actuator_kp_override(self, data_slice, kp: np.ndarray) -> None:
if not self._supports_position_actuator_gains:
raise NotImplementedError(
"Motrix actuator kp override is only available for all-position-actuator models"
)
for actuator in self._position_actuators:
# TODO(motrixsim#1384): drop the copy once strided NumPy views are accepted.
actuator.set_kp_override(data_slice, np.ascontiguousarray(kp[:, int(actuator.index)]))
def _set_position_actuator_kd_override(self, data_slice, kd: np.ndarray) -> None:
if not self._supports_position_actuator_gains:
raise NotImplementedError(
"Motrix actuator kd override is only available for all-position-actuator models"
)
for actuator in self._position_actuators:
# TODO(motrixsim#1384): drop the copy once strided NumPy views are accepted.
actuator.set_damping_override(
data_slice,
np.ascontiguousarray(kd[:, int(actuator.index)]),
)
[docs]
def get_actuator_gains(self) -> tuple[np.ndarray, np.ndarray]:
if not self._supports_position_actuator_gains:
raise NotImplementedError(
"Motrix actuator gains are only exposed for all-position-actuator models"
)
return self._default_actuator_kp.copy(), self._default_actuator_kd.copy()