Language Language English 简体中文

unisim.backend.mjwarp.backend

Host-compatibility implementation of the independent mjwarp backend.

mjwarp is not a MuJoCo backend mode. It uploads a CPU MuJoCo model to mujoco_warp and owns its own device data and host cache. The cache is refreshed exactly at explicit step/reset barriers; legacy getters only return views into that cache and therefore never trigger an implicit Warp .numpy transfer.

Classes

MjwarpBackend

Independent CUDA backend exposed through the host NumPy profile.

class unisim.backend.mjwarp.backend.MjwarpBackend

Bases: SimBackend

Independent CUDA backend exposed through the host NumPy profile.

State and control cross the host/device boundary only at explicit step/reset barriers with bounded, statically declared transfers. Reset DR writes per-world rows of cold-path-expanded model arrays in place and recomputes derived constants with the graded set_const* family; interval DR stages xfrc_applied pushes/body forces for the next step barrier or kicks root velocity through the reset upload+forward path. A registered pre-step control converter runs on the host before every physics substep (see set_pre_step_control); per-world gravity DR and native rendering remain fail-closed. Detached host snapshots support finite MuJoCo-based offline recording.

Parameters:

• scene (SceneCfg)

• num_envs (int)

• sim_dt (float)

• base_name (str | None)

• push_body_name (str | None)

• nconmax (int | None)

• njmax (int | None)

• add_body_sensors (bool)

• unexpected_kwargs (Any)

__init__(scene, num_envs, sim_dt, *, base_name=None, push_body_name=None, nconmax=None, njmax=None, add_body_sensors=False, **unexpected_kwargs)

Parameters:

• scene (SceneCfg)

• num_envs (int)

• sim_dt (float)

• base_name (str | None)

• push_body_name (str | None)

• nconmax (int | None)

• njmax (int | None)

• add_body_sensors (bool)

• unexpected_kwargs (Any)

property num_envs: int

Number of vectorized environments.

property model: Any

Return the backend-owned device model, never a MuJoCo backend model.

property num_actuators: int

Number of actuators.

property num_dof_vel: int

Number of joint velocity DoFs, excluding the floating base.

get_actuator_ctrl_range()

Return actuator control ranges.

Return type:

ndarray

Returns:

Array with shape (num_actuators, 2) and columns [low, high].

get_actuator_names()

Return actuator names in control-vector order on the cold path.

Return type:

tuple[str, ...]

get_actuator_joint_names()

Return each actuator’s target single-DoF joint in control-vector order.

Backends must fail closed when an actuator does not target exactly one hinge/slide joint. Manager action terms use this cold-path metadata to map community joint selectors onto the backend control vector without inspecting backend-private model objects.

Return type:

tuple[str, ...]

get_scene_model_file()

Return the materialized scene path for diagnostics, when available.

Return type:

str | None

get_keyframe_qpos(name)

Return the full qpos for a named keyframe, including the floating base.

Parameters:

name (str) – Keyframe name such as "stand" or "home".

Return type:

ndarray

Returns:

Array with shape (nq,).

get_default_qpos()

Return the backend/model default qpos through a stable contract.

Return type:

ndarray

get_default_dof_pos()

Return default joint positions in the same column order as get_dof_pos.

The returned array is detached, one-dimensional, and excludes floating root coordinates. Backends whose DoF view is actuator-indexed must use that same actuator-target order here.

Return type:

ndarray

get_init_qvel()

Return a zero-initialized qvel vector compatible with set_state.

Return type:

ndarray

Returns:

Zero-filled qvel array.

get_root_state_layout(root_body_name)

Resolve one body’s floating-root columns on the cold path.

Backends must verify that root_body_name owns a free/floating joint; fixed bodies and runtimes without body-to-root metadata fail closed. Name/model lookup is forbidden on reset and step hot paths, so callers cache either the returned layout or the unsupported result during scene materialization.

Parameters:

root_body_name (str)

Return type:

BackendRootStateLayout

get_body_ids(names)

Resolve body/link names to backend integer IDs.

Parameters:

names (Sequence[str]) – Body/link names.

Return type:

ndarray

Returns:

int32 array with shape (len(names),).

Raises:

ValueError – If any name is not found.

get_geom_id(name)

Resolve one geom name through the backend contract.

Parameters:

name (str)

Return type:

int

get_geom_size(name)

Return one geom size vector through the backend contract.

Parameters:

name (str)

Return type:

ndarray

get_geom_sizes()

Return default geometry sizes, shape (ngeom, 3).

Return type:

ndarray

get_geom_solref()

Return default contact reference parameters, shape (ngeom, 2).

Return type:

ndarray

get_geom_solimp()

Return default contact impedance parameters, shape (ngeom, 5).

Return type:

ndarray

get_dof_damping()

Return default joint damping, shape (nv,).

Return type:

ndarray

get_dof_frictionloss()

Return default joint friction loss, shape (nv,).

Return type:

ndarray

bind_mocap_pose(body_name)

Resolve a mocap body once; unavailable capabilities fail at binding.

Parameters:

body_name (str)

Return type:

BackendMocapPoseBinding

get_body_subtree_ids(root_body_id)

Return body ids in the subtree rooted at root_body_id.

Parameters:

root_body_id (int)

Return type:

ndarray

get_geom_names()

Return backend geom names in backend id order.

Return type:

tuple[str, ...]

get_geom_body_ids()

Return the owning body id for each geom.

Return type:

ndarray

get_geom_contact_masks()

Return per-geom contact type and affinity masks.

Return type:

tuple[ndarray, ndarray]

get_geom_friction()

Return the backend geom-friction table.

Return type:

ndarray

get_gravity()

Return the backend gravity vector.

Return type:

ndarray

get_body_mass()

Return the backend body-mass table.

Return type:

ndarray

get_body_ipos()

Return the backend body inertial-position table.

Return type:

ndarray

get_dof_armature()

Return the backend dof-armature table.

Return type:

ndarray

get_motion_body_ids(names)

Resolve backend-native body IDs used by motion datasets.

Parameters:

names (Sequence[str])

Return type:

ndarray

get_joint_range()

Return joint position limits, excluding the floating base.

Return type:

ndarray | None

Returns:

Array with shape (num_dof, 2) and columns [low, high], or None when the backend does not expose limits.

get_site_ids(names)

Resolve site names to integer ID arrays.

Parameters:

names (Sequence[str]) – Site names.

Return type:

ndarray

Returns:

int32 ID array with shape (len(names),).

get_joint_dof_indices(names)

Resolve named joint qvel coordinates on the cold metadata path.

Parameters:

names (Sequence[str])

Return type:

ndarray

get_joint_dof_pos_indices(names)

Resolve named single-DoF qpos coordinates excluding the free root.

Parameters:

names (Sequence[str])

Return type:

ndarray

get_joint_dof_vel_indices(names)

Resolve named joint qvel coordinates excluding the free root.

Parameters:

names (Sequence[str])

Return type:

ndarray

get_joint_state_qpos_indices(names)

Resolve named joints to full reset qpos columns.

Parameters:

names (Sequence[str])

Return type:

ndarray

get_joint_state_qvel_indices(names)

Resolve named joints to full reset qvel columns.

Parameters:

names (Sequence[str])

Return type:

ndarray

get_actuator_gains()

Expose immutable model defaults; this does not advertise gain DR support.

Return type:

tuple[ndarray, ndarray]

set_pre_step_control(fn)

Register or clear the env-owned per-substep control converter.

Semantics match the MuJoCo backend: the callback runs once before every physics substep with the host qpos/qvel cache refreshed to the substep-start state, and its return value becomes that substep’s device ctrl. Each invocation costs one explicit device round trip, so the callback path intentionally uses eager kernel launches instead of replaying the captured step graph. Passing None restores the direct control path.

Parameters:

fn (Callable[[Any, ndarray], ndarray] | None)

Return type:

None

step(ctrl, nsteps=1)

Advance physics.

Parameters:

• ctrl (ndarray) – Control input with shape (num_envs, nu).

• nsteps (int) – Number of physics substeps.

Return type:

dict[str, dict[str, float]]

Returns:

Optional dictionary. Backends may include a "timing" key with per-phase timings in milliseconds.

set_state(env_indices, qpos, qvel, randomization=None)

Set physics state for selected environments.

Parameters:

• env_indices (ndarray) – Environment indices.

• qpos (ndarray) – Position state. Free-root columns exposed by get_root_state_layout() use world xyz and wxyz quaternion.

• qvel (ndarray) – Velocity state. Free-root columns exposed by get_root_state_layout() use world linear velocity and body-frame angular velocity.

• randomization (ResetRandomizationPayload | None) – Optional backend randomization payload.

Return type:

dict[str, dict[str, float]]

Returns:

Optional dictionary. Backends MAY include a "timing" key with per-substep timings in milliseconds (e.g. set_state_mask_ms, set_state_data_slice_ms, …). Callers MUST treat None or missing keys as “not reported” — the outer wall-clock measurement in DomainRandomizationManager.reset (dr_reset_set_state_ms) remains authoritative for total set_state time.

get_dr_capabilities()

Advertise the per-world model mutation set validated by effect tests.

Return type:

DomainRandomizationCapabilities

push_robots(force_range)

Sample one world-frame push force per env and stage it for the next step.

Parameters:

force_range (Sequence[float] | ndarray)

Return type:

None

apply_body_force(body_ids, force, torque=None)

Accumulate world-frame forces on the staged xfrc_applied rows.

Parameters:

• body_ids (ndarray)

• force (ndarray)

• torque (ndarray | None)

Return type:

None

materialize()

Resources are fully materialized during the constructor cold path.

Return type:

None

get_play_capabilities()

Return backend-native play/render capabilities.

Return type:

BackendPlayCapabilities

resolve_play_render_plan(*, play_render_mode, play_steps, output_video)

Resolve high-level playback mode into backend-owned render parameters.

Parameters:

• play_render_mode (str | None)

• play_steps (int | None)

• output_video (str | PathLike[str] | None)

Return type:

BackendPlayRenderPlan

run_playback(*, env, initialize, step, num_steps, output_video=None, render_spacing=None, render_offset_mode=None, headless=None, record_video=None, frame_state_getter=None, camera_kwargs=None, debug_overlay_getter=None, on_frame=None)

Execute backend-owned playback for an env wrapper.

camera_kwargs is normalized into CameraCfg at this boundary; unknown mapping keys fail closed with an error naming them.

debug_overlay_getter is an optional per-frame callback returning a sequence with one entry per environment (len == num_envs); each entry is that env’s sequence of DebugPrimitive (None or empty marks an env without overlay) and returning None disables overlays for the frame. Primitive poses are env-local; the renderer applies grid offsets when composing multiple envs. Backends whose get_play_capabilities().supports_debug_overlay is False fail closed with NotImplementedError when this is not None. On the interactive rendering path only backends whose supports_interactive_debug_overlay is True consume it; the others fail closed with NotImplementedError.

on_frame is an optional per-frame video hook called by offline render pipelines before encoding: it receives (frame_index, frame) with the frame an (H, W, 3) uint8 array, and returns a replacement frame of the same shape/dtype or None to keep the original. Backends rendering through a native (non-offline) renderer fail closed with NotImplementedError when this is not None.

Known boundary: env is the owning env wrapper, not a physics-layer concept. Current playback implementations read env-level configuration (e.g. cfg.scene, cfg.ctrl_dt, cfg.render_spacing) and env-owned playback helpers (get_playback_model, get_physics_state_snapshot) that have no backend-native equivalent yet. The parameter stays on this contract until playback asset/config resolution moves onto backend-owned metadata; backends must only use it on the cold playback path.

Parameters:

• env (Any)

• initialize (Any)

• step (Any)

• num_steps (int | None)

• output_video (str | PathLike[str] | None)

• render_spacing (float | None)

• render_offset_mode (str | None)

• headless (bool | None)

• record_video (bool | None)

• frame_state_getter (Any)

• camera_kwargs (CameraCfg | Mapping[str, Any] | None)

• debug_overlay_getter (DebugOverlayGetter | None)

• on_frame (Any)

Return type:

str | None

get_physics_state()

Return a physics snapshot suitable for offline playback/video export.

Rows use the [time, qpos, qvel] layout; backends whose model has mocap bodies append [mocap_pos(nmocap*3), mocap_quat(nmocap*4)] so offline rendering can replay mocap-driven geometry at its recorded pose.

Return type:

ndarray

get_playback_mocap_state(env_index=0)

Return copied mocap pose arrays for detached visual playback.

Parameters:

env_index (int)

Return type:

tuple[ndarray, ndarray]

get_playback_model(env_index=None)

Return the playback model for a specific env when variants exist.

Parameters:

env_index (int | None) – Optional vectorized environment index.

Return type:

str

Returns:

The backend model object used by playback tooling.

get_base_pos()

Return base position in the world frame.

Return type:

ndarray

Returns:

(num_envs, 3)

get_base_quat()

Return base quaternion in the world frame as wxyz.

Return type:

ndarray

Returns:

(num_envs, 4)

get_base_lin_vel()

Return base linear velocity in the world frame.

This is the first three dimensions of generalized velocity qvel, expressed in world coordinates.

Return type:

ndarray

Returns:

(num_envs, 3)

get_base_ang_vel()

Return base angular velocity in the world frame.

This is dimensions 3-5 of generalized velocity qvel, expressed in world coordinates. It differs from gyro readings: gyro sensors report angular velocity components in the body/sensor local frame, while this contract returns world-frame values. Use the matching sensor contract when body-frame angular velocity is required.

Return type:

ndarray

Returns:

(num_envs, 3)

get_dof_pos()

Return joint positions, excluding the base.

Return type:

ndarray

Returns:

(num_envs, num_dof)

get_dof_vel()

Return joint velocities, excluding the base.

Return type:

ndarray

Returns:

(num_envs, num_dof)

get_body_pos_w(body_ids)

Return selected body positions in the world frame.

Parameters:

body_ids (ndarray) – Body ID array.

Return type:

ndarray

Returns:

(num_envs, len(body_ids), 3)

get_body_quat_w(body_ids)

Return selected body quaternions in the world frame as wxyz.

Parameters:

body_ids (ndarray) – Body ID array.

Return type:

ndarray

Returns:

(num_envs, len(body_ids), 4)

get_body_pose_w_rows(env_ids, body_ids)

Gather world-frame body pose for selected environments only.

Parameters:

• env_ids (ndarray)

• body_ids (ndarray)

Return type:

tuple[ndarray, ndarray]

get_body_lin_vel_w(body_ids)

Return selected body linear velocities in the world frame.

Parameters:

body_ids (ndarray) – Body ID array.

Return type:

ndarray

Returns:

(num_envs, len(body_ids), 3)

get_body_ang_vel_w(body_ids)

Return selected body angular velocities in the world frame.

Parameters:

body_ids (ndarray) – Body ID array.

Return type:

ndarray

Returns:

(num_envs, len(body_ids), 3)

get_body_lin_vel_w_rows(env_ids, body_ids)

Gather world-frame body linear velocity for selected rows.

Parameters:

• env_ids (ndarray)

• body_ids (ndarray)

Return type:

ndarray

get_body_ang_vel_w_rows(env_ids, body_ids)

Gather world-frame body angular velocity for selected rows.

Parameters:

• env_ids (ndarray)

• body_ids (ndarray)

Return type:

ndarray

copy_body_state_w(body_ids, out_pos, out_quat, out_lin_vel, out_ang_vel)

Copy selected world-frame body state into caller-owned buffers.

Parameters:

• body_ids (ndarray)

• out_pos (ndarray)

• out_quat (ndarray)

• out_lin_vel (ndarray)

• out_ang_vel (ndarray)

Return type:

tuple[ndarray, ndarray, ndarray, ndarray]

get_body_pos_b(body_ids)

Return selected body positions in the baselink frame.

Parameters:

body_ids (ndarray) – Body ID array.

Return type:

ndarray

Returns:

(num_envs, len(body_ids), 3)

get_body_quat_b(body_ids)

Return selected body quaternions in the baselink frame as wxyz.

Parameters:

body_ids (ndarray) – Body ID array.

Return type:

ndarray

Returns:

(num_envs, len(body_ids), 4)

get_body_lin_vel_b(body_ids)

Return selected body linear velocities expressed in each body’s own frame.

The value is the body’s world-frame velocity rotated by the inverse of the body’s world-frame orientation, i.e. quat_apply_inverse(quat_w, lin_vel_w) (mjlab/Isaac-style analytical definition). It is well-defined for every body — including the root body — and must NOT be implemented as the motion relative to the baselink frame (which degenerates to zero for the root body).

Parameters:

body_ids (ndarray) – Body ID array.

Return type:

ndarray

Returns:

(num_envs, len(body_ids), 3)

get_body_ang_vel_b(body_ids)

Return selected body angular velocities expressed in each body’s own frame.

The value is the body’s world-frame angular velocity rotated by the inverse of the body’s world-frame orientation, i.e. quat_apply_inverse(quat_w, ang_vel_w) (mjlab/Isaac-style analytical definition). It is well-defined for every body — including the root body — and must NOT be implemented as the motion relative to the baselink frame (which degenerates to zero for the root body).

Parameters:

body_ids (ndarray) – Body ID array.

Return type:

ndarray

Returns:

(num_envs, len(body_ids), 3)

get_sensor_data(name)

Return sensor data.

Parameters:

name (str) – Sensor name.

Return type:

ndarray

Returns:

Sensor data array.