Language

unisim.backend.drake.backend¶

UniSim adapter for the drake-uni batch runtime.

UniLab owns task logic, reset sampling, named sensor views, and training flow. DrakeUni owns Drake model construction, batched stepping, and raw sensor evaluation. This module translates the SimBackend contract into DrakeUni runtime calls and keeps UniLab’s cached state/sensor views synchronized.

Functions

ensure_drake_batch_available()

Report whether the DrakeUni batch extension can be used.

Classes

DrakeBackend

UniLab SimBackend implementation backed by DrakeUni batch runtime.

class unisim.backend.drake.backend.DrakeBackend[source]¶

Bases: SimBackend

UniLab SimBackend implementation backed by DrakeUni batch runtime.

The backend keeps the public UniLab API stable while delegating model construction, integration, and raw sensor evaluation to DrakeUni.

Parameters:

scene (SceneCfg)
num_envs (int)
sim_dt (float)
drake_backend_mode (str)
nthread (int)

backend_type: str = 'drake'¶

__init__(scene, num_envs, sim_dt, *, drake_backend_mode='batch', nthread=0)[source]¶

Parameters:

scene (SceneCfg)
num_envs (int)
sim_dt (float)
drake_backend_mode (str)
nthread (int)

property scene_model_file: str¶

property num_envs: int¶: Number of vectorized environments.

property nthread: int¶

property model: _DrakeUniModelView¶: Underlying physics 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()[source]¶

Return actuator control ranges.

Return type:: ndarray
Returns:: Array with shape (num_actuators, 2) and columns [low, high].

get_actuator_names()[source]¶

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

Return type:: tuple[str, ...]

get_actuator_joint_names()[source]¶

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()[source]¶

Return the materialized scene path for diagnostics, when available.

Return type:: str | None

get_joint_range()[source]¶

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_keyframe_qpos(name)[source]¶

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()[source]¶

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

Return type:: ndarray

get_default_dof_pos()[source]¶

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()[source]¶

Return a zero-initialized qvel vector compatible with set_state.

Return type:: ndarray
Returns:: Zero-filled qvel array.

get_actuator_gains()[source]¶

Return per-joint (kp, kd) arrays from the backend model.

Return type:: tuple[ndarray, ndarray]

get_body_ids(names)[source]¶

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_motion_body_ids(names)[source]¶

Resolve backend-native body IDs used by motion datasets.

Parameters:: names (Sequence[str])
Return type:: ndarray

get_site_ids(names)[source]¶

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)[source]¶

Resolve joint names to DoF indices in velocity space (qvel).

Parameters:: names (Sequence[str]) – Joint names.
Return type:: ndarray
Returns:: int32 index array with shape (len(names),) relative to the qvel start.

get_joint_dof_pos_indices(names)[source]¶

Resolve joint names to DoF indices in position space (qpos).

Only single-DoF joints are supported; free joints are excluded.

Parameters:: names (Sequence[str]) – Joint names.
Return type:: ndarray
Returns:: int32 index array with shape (len(names),) relative to the joint section of qpos.

get_joint_dof_vel_indices(names)[source]¶

Resolve joint names to DoF indices in velocity space (qvel).

Parameters:: names (Sequence[str]) – Joint names.
Return type:: ndarray
Returns:: int32 index array with shape (len(names),) relative to the joint section start.

get_joint_state_qpos_indices(names)[source]¶

Resolve single-DoF joints to full set_state qpos columns.

Unlike get_joint_dof_pos_indices(), these indices address the complete qpos vector accepted by set_state(), including any root coordinates. Manager reset transactions resolve them on the cold path.

Parameters:: names (Sequence[str])
Return type:: ndarray

get_joint_state_qvel_indices(names)[source]¶

Resolve single-DoF joints to full set_state qvel columns.

Parameters:: names (Sequence[str])
Return type:: ndarray

step(ctrl, nsteps=1)[source]¶

Advance physics.

Parameters:

ctrl (ndarray) – Control input with shape (num_envs, nu).
nsteps (int) – Number of physics substeps.

Return type:

dict | None

Returns:

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

set_state(env_indices, qpos, qvel, randomization=None)[source]¶

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:

None

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()[source]¶

Return supported domain-randomization capabilities for this backend.

Return type:: DomainRandomizationCapabilities

apply_interval_randomization(plan)[source]¶

Apply a scheduled interval randomization plan.

Generic dispatch: each op yielded by plan.iter_ops() is validated against the builtin term specs (custom terms pass through) and routed to the backend-owned handler table returned by _interval_term_handlers(). A term without a handler fails closed with NotImplementedError naming the backend class and the term. Backends that need per-plan prologue/epilogue semantics (for example clearing staged external forces before the ops accumulate) keep a thin override that calls this base implementation.

Parameters:: plan (IntervalRandomizationPlan)
Return type:: None

get_play_capabilities()[source]¶

Return backend-native play/render capabilities.

Return type:: BackendPlayCapabilities

resolve_play_render_plan(*, play_render_mode, play_steps, output_video)[source]¶

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)[source]¶

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 (Callable[[], Any])
step (Callable[[Any], 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 (Callable[[], np.ndarray] | None)
camera_kwargs (CameraCfg | Mapping[str, Any] | None)
debug_overlay_getter (DebugOverlayGetter | None)
on_frame (Callable[[int, np.ndarray], np.ndarray | None] | None)

Return type:

str | None

init_renderer(spacing=1.0, *, offset_mode='grid', headless=False, capture=False, width=1280, height=720, camera_kwargs=None)[source]¶

Initialize a backend-native renderer.

headless controls whether a native window is opened. capture controls whether capture_video_frame is valid for the renderer. camera_kwargs is normalized into CameraCfg at this boundary; unknown mapping keys fail closed with an error naming them.

Parameters:

spacing (float)
offset_mode (str)
headless (bool)
capture (bool)
width (int)
height (int)
camera_kwargs (dict[str, Any] | None)

Return type:

None

render()[source]¶

Render one frame through a backend-native interactive renderer.

Raises:: RenderClosedError – If the user closed the render window.
Return type:: None

capture_video_frame()[source]¶

Capture one RGB frame through a backend-native renderer.

Raises:: RenderClosedError – If the user closed the render window.
Return type:: ndarray

get_physics_state()[source]¶

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_model(env_index=None)[source]¶

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.

diagnostics()[source]¶

Return type:: Any

apply_body_force(body_ids, force, torque=None)[source]¶

Apply a world-frame force (and optional torque) to bodies for the upcoming step.

Parameters:

body_ids (ndarray) – Body ids whose external forces should be perturbed.
force (ndarray) – Force values with shape (num_envs, len(body_ids), 3).
torque (ndarray | None) – Optional world-frame torque values with the same shape. Backends without a torque channel must fail closed when this is not None.

Return type:

None

Returns:

None. Backends that support this mutate their pending simulation state.

get_base_pos()[source]¶

Return base position in the world frame.

Return type:: ndarray
Returns:: (num_envs, 3)

get_base_quat()[source]¶

Return base quaternion in the world frame as wxyz.

Return type:: ndarray
Returns:: (num_envs, 4)

get_base_lin_vel()[source]¶

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()[source]¶

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()[source]¶

Return joint positions, excluding the base.

Return type:: ndarray
Returns:: (num_envs, num_dof)

get_dof_vel()[source]¶

Return joint velocities, excluding the base.

Return type:: ndarray
Returns:: (num_envs, num_dof)

get_body_pos_w(body_ids)[source]¶

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)[source]¶

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_lin_vel_w(body_ids)[source]¶

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)[source]¶

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_pos_b(body_ids)[source]¶

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)[source]¶

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)[source]¶

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)[source]¶

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)[source]¶

Return sensor data.

Parameters:: name (str) – Sensor name.
Return type:: ndarray
Returns:: Sensor data array.

unisim.backend.drake.backend.ensure_drake_batch_available()[source]¶

Report whether the DrakeUni batch extension can be used.

Return type:: tuple[bool, ImportError | None]