unisim.backend.mjwarp.backend.MjwarpBackend¶
- class unisim.backend.mjwarp.backend.MjwarpBackend[source]¶
Bases:
SimBackendIndependent 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 stagesxfrc_appliedpushes/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 (seeset_pre_step_control) and may additionally return a per-substep body wrench that composes with the staged interval wrenches; native rendering remains fail-closed. Detached host snapshots support finite MuJoCo-based offline recording.- Parameters:
Methods
__init__(scene, num_envs, sim_dt, *[, ...])apply_body_force(body_ids, force[, torque])Accumulate a world-frame wrench on the staged
xfrc_appliedrows.Apply a scheduled interval randomization plan.
bind_height_scanner(*, hfield_geom_id, ...)Create a reusable height-field scanner on the init/cold path.
bind_host_bridge_io(spec)Compile a persistent transfer plan for a declared host bridge.
bind_mocap_pose(body_name)Resolve a mocap body once; unavailable capabilities fail at binding.
bind_packed_reset(*[, device])Begin a device-resident transaction with one scatter commit boundary.
bind_renderer([spacing, offset_mode, ...])Initialize a backend-native renderer.
bind_sensor_data(names)Materialize a validated view over named tensor sensors.
Capture one RGB frame through a backend-native renderer.
Release cold-path scene artifacts owned by the backend.
close()Return actuator control ranges.
Expose immutable model defaults; this does not advertise gain DR support.
Return each actuator's target single-DoF joint in control-vector order.
Return actuator names in control-vector order on the cold path.
get_body_id(name)Resolve one body/link name through the backend contract.
get_body_ids(names)Resolve body/link names to backend integer IDs.
get_body_ipos([env_ids])Return body center-of-mass offsets in each body's local frame, in meters.
Return the backend body-mass table.
get_body_subtree_ids(root_body_id)Return body ids in the subtree rooted at
root_body_id.get_capabilities(*[, profile])Query semantic declarations plus authoritative DR/play/variant support.
Return default joint positions in canonical non-root qpos order.
Return default geometry sizes, shape (ngeom, 3).
Return the backend/model default qpos through a stable contract.
Return the backend dof-armature table.
Return default joint damping, shape (nv,).
Return default joint friction loss, shape (nv,).
Advertise the per-world model mutation set validated by effect tests.
get_entity_default_state(entity[, env_ids])Detached construction/keyframe defaults in selected environment order.
Return materialized physical entity names in the frozen public order.
get_entity_state(entity)Return detached root_pose/root_velocity/joint_positions/joint_velocities.
Return the owning body id for each geom.
Return per-geom contact type and affinity masks.
Return the backend geom-friction table.
get_geom_id(name)Resolve one geom name through the backend contract.
Return backend geom names in backend id order.
get_geom_size(name)Return one geom size vector in backend geom-name order.
Return default contact impedance parameters, shape (ngeom, 5).
Return default contact reference parameters, shape (ngeom, 2).
Return the backend gravity vector.
Return a detached construction/materialization configuration snapshot.
Return a zero-initialized qvel vector compatible with
set_state_tensor.get_joint_dof_indices(names)Resolve named joint qvel coordinates on the cold metadata path.
get_joint_dof_pos_indices(names)Resolve named single-DoF qpos coordinates excluding the free root.
get_joint_dof_vel_indices(names)Resolve named joint qvel coordinates excluding the free root.
get_joint_range(*[, names])Return joint position limits, excluding the floating base.
get_joint_state_qpos_indices(names)Resolve named joints to full reset qpos columns.
get_joint_state_qvel_indices(names)Resolve named joints to full reset qvel columns.
get_keyframe_qpos(name)Return the full qpos for a named keyframe, including the floating base.
get_motion_body_ids(names)Resolve body IDs used by motion datasets.
Return a physics snapshot suitable for offline playback/video export.
Return the snapshot layout, including the mocap tail when present.
Return backend-native play/render capabilities.
get_playback_mocap_state([env_index])Return copied mocap pose arrays for detached visual playback.
get_playback_model([env_index])Return the playback model for a specific env when variants exist.
Return canonical qpos/qvel widths for packed tensor reset layout.
get_reset_term_default(term)Return canonical or fixed-variant authoritative reset defaults.
get_root_state_layout(root_body_name)Resolve one body's floating-root columns on the cold path.
Return the immutable materialized public addresses, never native handles.
Return the materialized scene path for diagnostics, when available.
Return the scene visual model file on the cold path, when available.
Return the complete public named-sensor inventory and widths.
Return the complete public named-sensor namespace.
get_sensor_view(name[, device])Return a live named-sensor view without exposing sensor offsets.
get_site_ids(names)Resolve site names to integer ID arrays.
get_site_jacobian_w(site_id, dof_indices)Compute world-frame Jacobians for one site and selected DoF columns.
get_state_views([fields, device])Return live Torch views into stable MJWarp device fields.
Return fail-closed tensor methods and negotiable state fields.
Declare the optional tensor lifecycle without discovering SDKs.
Expose whether fixed-address MJWarp graph replay is active.
Return backend-materialized terrain metadata on the cold path.
get_tracked_body_views([body_names, device])Return four ordered tracked-body blocks from one stable device read.
Finalize cold-path backend resources before reset/step.
push_robots(force_range)Stage caller-owned device push forces for the next tensor step.
render()Render one frame through a backend-native interactive renderer.
reset_entities(request)Validate the complete selected-entity request, then commit it once.
resolve_play_render_plan(*, ...)Resolve high-level playback mode into backend-owned render parameters.
run_playback(*, env, initialize, step, num_steps)Execute backend-owned playback for an env wrapper.
set_mocap_pose_tensor(body_name, ...)Write selected mocap poses through the declared tensor lifecycle.
set_physics_state(state)Restore a snapshot produced by
get_physics_state.Register an env-owned policy-control to physics-control converter.
set_state_tensor(env_indices, qpos, qvel[, ...])Commit selected generalized state with device-resident scatter kernels.
set_state_tensor_mask(reset_mask, qpos, qvel)Commit full-width tensors selected by a backend-published bool mask.
step_tensor(ctrl[, nsteps])Advance MJWarp physics without a control or state host transfer.
Attributes
Coarse capability labels for clients that need a cheap feature check.
Return the backend-owned device model, never a MuJoCo backend model.
Number of actuators.
Number of joint velocity DoFs, excluding the floating base.
Number of vectorized environments.
- __init__(scene, num_envs, sim_dt, *, base_name=None, push_body_name=None, nconmax=None, njmax=None, add_body_sensors=False, **unexpected_kwargs)[source]¶
- get_scene_layout()[source]¶
Return the immutable materialized public addresses, never native handles.
- Return type:
CompiledSceneLayout
- get_entity_default_state(entity, env_ids=None)[source]¶
Detached construction/keyframe defaults in selected environment order.
Fields and frames match get_entity_state. Defaults follow immutable variant assignment and never reflect current reset-time randomization. This query does not reset, step or reparse the scene.
- get_entity_state(entity)[source]¶
Return detached root_pose/root_velocity/joint_positions/joint_velocities.
Root fields follow EntityStatePatch’s link-origin/world-frame contract. Non-root joint fields use the frozen entity joint order and native generalized widths. Unavailable fields must fail, never return old data as if current. State freshness follows the adapter’s declared profile.
- reset_entities(request)[source]¶
Validate the complete selected-entity request, then commit it once.
Validation failure leaves all state unchanged. A native partial commit failure must fault the backend unless the adapter can actually roll back. Unselected entities/environments and fixed identity remain unchanged.
- Parameters:
request (
SceneResetRequest)- Return type:
- cleanup_scene_assets()[source]¶
Release cold-path scene artifacts owned by the backend.
- Return type:
- get_actuator_ctrl_range()[source]¶
Return actuator control ranges.
- Return type:
- Returns:
Array with shape
(num_actuators, 2)and columns[low, high].
- 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.
- get_keyframe_qpos(name)[source]¶
Return the full qpos for a named keyframe, including the floating base.
- get_default_qpos()[source]¶
Return the backend/model default qpos through a stable contract.
- Return type:
- get_default_dof_pos()[source]¶
Return default joint positions in canonical non-root qpos order.
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:
- get_init_qvel()[source]¶
Return a zero-initialized qvel vector compatible with
set_state_tensor.- Return type:
- Returns:
Zero-filled qvel array.
- get_tensor_execution()[source]¶
Declare the optional tensor lifecycle without discovering SDKs.
- Return type:
TensorExecution
- get_tensor_capabilities()[source]¶
Return fail-closed tensor methods and negotiable state fields.
- Return type:
TensorLifecycleCapabilities
- bind_packed_reset(*, device=None)[source]¶
Begin a device-resident transaction with one scatter commit boundary.
- get_public_state_widths()[source]¶
Return canonical qpos/qvel widths for packed tensor reset layout.
- Return type:
PublicStateWidths
- get_tensor_runtime_diagnostics()[source]¶
Expose whether fixed-address MJWarp graph replay is active.
- get_state_views(fields=None, device=None)[source]¶
Return live Torch views into stable MJWarp device fields.
- get_sensor_view(name, device=None)[source]¶
Return a live named-sensor view without exposing sensor offsets.
- get_tracked_body_views(body_names=None, device=None)[source]¶
Return four ordered tracked-body blocks from one stable device read.
- get_root_state_layout(root_body_name)[source]¶
Resolve one body’s floating-root columns on the cold path.
Backends must verify that
root_body_nameowns 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:
RootStateLayout
- get_body_ids(names)[source]¶
Resolve body/link names to backend integer IDs.
- Parameters:
- Return type:
- Returns:
int32array with shape(len(names),).- Raises:
ValueError – If any name is not found.
- get_geom_size(name)[source]¶
Return one geom size vector in backend geom-name order.
get_default_geom_sizesreturns the complete table of default sizes, while this selector resolves one geom by name.
- get_geom_solref()[source]¶
Return default contact reference parameters, shape (ngeom, 2).
- Return type:
- get_geom_solimp()[source]¶
Return default contact impedance parameters, shape (ngeom, 5).
- Return type:
- bind_mocap_pose(body_name)[source]¶
Resolve a mocap body once; unavailable capabilities fail at binding.
- Parameters:
body_name (
str)- Return type:
MocapPoseBinding
- get_body_ipos(env_ids=None)[source]¶
Return body center-of-mass offsets in each body’s local frame, in meters.
body_ipos[b]is the position of bodyb’s center of mass in that body’s own local coordinate frame; it is neither an inertia tensor nor a whole-subtree COM.Without
env_idsthe canonical model default table is returned with shape(nbody, 3)in backend body-id order. It never changes with reset randomization; fixed-variant backends expose their per-environment default baselines throughget_reset_term_default()instead.With
env_idsthe current effective per-environment values are returned with shape(len(env_ids), nbody, 3), inenv_idsorder. The values reflect every reset randomization applied so far, including composition withbase_com_offset; an environment untouched by a partial reset keeps its previous values. Backends that do not track per-environment inertial offsets fail closed withNotImplementedErrorfor this form.
- get_motion_body_ids(names)[source]¶
Resolve body IDs used by motion datasets.
Motion datasets are generated from MuJoCo, so the returned ids follow the MJCF body order with
worldbodyas id 0, regardless of the backend-native indexing used byget_body_ids.
- get_joint_range(*, names=None)[source]¶
Return joint position limits, excluding the floating base.
- Parameters:
names (
Sequence[str] |None) – Optional joint names in requested return order. Adapters that support named scalar joints resolve the model mapping internally.- Return type:
- Returns:
Array with shape
(num_dof, 2)and columns[low, high], orNonewhen the backend does not expose limits.When
namesis provided by a supporting adapter, the result has shape(len(names), 2). Hinge limits are radians, slide limits are meters, and joints without enabled limits return infinities.
- get_joint_dof_indices(names)[source]¶
Resolve named joint qvel coordinates on the cold metadata path.
- get_joint_dof_pos_indices(names)[source]¶
Resolve named single-DoF qpos coordinates excluding the free root.
- get_joint_dof_vel_indices(names)[source]¶
Resolve named joint qvel coordinates excluding the free root.
- get_actuator_gains()[source]¶
Expose immutable model defaults; this does not advertise gain DR support.
- step_tensor(ctrl, nsteps=1)[source]¶
Advance MJWarp physics without a control or state host transfer.
- set_state_tensor(env_indices, qpos, qvel, randomization=None)[source]¶
Commit selected generalized state with device-resident scatter kernels.
- set_state_tensor_mask(reset_mask, qpos, qvel, randomization=None, *, full_refresh=False)[source]¶
Commit full-width tensors selected by a backend-published bool mask.
The public full-width state payload stays fixed-shape. The backend owns mask validation and reset-graph publication, so the caller never probes private row compaction or CUDA graph state.
- get_reset_term_default(term)[source]¶
Return canonical or fixed-variant authoritative reset defaults.
- get_dr_capabilities()[source]¶
Advertise the per-world model mutation set validated by effect tests.
- 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(). Op payloads are torch-first: Torch tensors orTensorIntervalPayloadvalues produced by the caller; host adapters convert tensors at their private boundary throughunisim.dr._host, and a missing producer-sampled payload fails closed. A term without a handler fails closed withNotImplementedErrornaming 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:
- push_robots(force_range)[source]¶
Stage caller-owned device push forces for the next tensor step.
The producer owns random generation on the backend Torch device. This boundary receives final full-width forces, validates them on their device, and scatters directly into the live
xfrc_appliedfield; it does not sample on the host, publish the host cache, or mutate the host wrench mirror.- Parameters:
force_range (
TensorIntervalPayload)- Return type:
- apply_body_force(body_ids, force, torque=None)[source]¶
Accumulate a world-frame wrench on the staged
xfrc_appliedrows.The force acts at each target body’s center of mass and the torque is about that center, matching MuJoCo
xfrc_appliedsemantics. The staged wrench applies to every substep of the nextstep()call and is cleared after it; repeated calls accumulate until consumed.
- get_play_capabilities()[source]¶
Return backend-native play/render capabilities.
- Return type:
PlayCapabilities
- resolve_play_render_plan(*, play_render_mode, play_steps, output_video)[source]¶
Resolve high-level playback mode into backend-owned render parameters.
- 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_kwargsis normalized intoCameraCfgat this boundary; unknown mapping keys fail closed with an error naming them.debug_overlay_getteris an optional per-frame callback returning a sequence with one entry per environment (len == num_envs); each entry is that env’s sequence ofDebugPrimitive(Noneor empty marks an env without overlay) and returningNonedisables overlays for the frame. Primitive poses are env-local; the renderer applies grid offsets when composing multiple envs. Backends whoseget_play_capabilities().supports_debug_overlayis False fail closed withNotImplementedErrorwhen this is notNone. On the interactive rendering path only backends whosesupports_interactive_debug_overlayis True consume it; the others fail closed withNotImplementedError.on_frameis 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 orNoneto keep the original. Backends rendering through a native (non-offline) renderer fail closed withNotImplementedErrorwhen this is notNone.Known boundary:
envis 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_layout()[source]¶
Return the snapshot layout, including the mocap tail when present.
- Return type:
PhysicsStateLayout
- 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. Consumers must decode snapshots throughget_physics_state_layout().split_staterather than hardcoding column slices.- Return type:
- get_playback_mocap_state(env_index=0)[source]¶
Return copied mocap pose arrays for detached visual playback.
- get_playback_model(env_index=None)[source]¶
Return the playback model for a specific env when variants exist.
- bind_height_scanner(*, hfield_geom_id, offsets, frame_body_id, alignment='yaw', output='height')¶
Create a reusable height-field scanner on the init/cold path.
Backends that support height-field terrain scan must override this method.
- bind_host_bridge_io(spec)¶
Compile a persistent transfer plan for a declared host bridge.
- Parameters:
spec (
TensorIOSpec)- Return type:
HostBridgeTransferPlan
- bind_renderer(spacing=1.0, *, offset_mode='grid', headless=False, capture=False, width=1280, height=720, camera_kwargs=None)¶
Initialize a backend-native renderer.
headlesscontrols whether a native window is opened.capturecontrols whethercapture_video_frameis valid for the renderer.camera_kwargsis normalized intoCameraCfgat this boundary; unknown mapping keys fail closed with an error naming them.
- bind_sensor_data(names)¶
Materialize a validated view over named tensor sensors.
This cold-path binding records the flattened public view width and delegates execution to the backend-owned tensor sensor reader. It never falls back to a NumPy named-sensor getter.
- property capability_flags¶
Coarse capability labels for clients that need a cheap feature check.
The detailed contract is expressed by the methods on this class. The labels remain useful for benchmark/conformance metadata and are derived from the mandatory lifecycle methods rather than maintained separately by every adapter.
- capture_video_frame()¶
Capture one RGB frame through a backend-native renderer.
- Raises:
RenderClosedError – If the user closed the render window.
- Return type:
- get_body_id(name)¶
Resolve one body/link name through the backend contract.
- get_capabilities(*, profile='default')¶
Query semantic declarations plus authoritative DR/play/variant support.
This cold-path query does not discover SDKs or verify runtime behavior. Materialized configuration readback is exposed by
get_import_report.- Parameters:
profile (
str)- Return type:
CapabilityReport
- get_import_report()¶
Return a detached construction/materialization configuration snapshot.
This cached report never parses assets, starts workers, or reflects reset-time randomization. Use current-property queries for live values.
- Return type:
ImportReport
- get_scene_visual_model_file()¶
Return the scene visual model file on the cold path, when available.
Backends without a separate visual scene model return
None.
- get_site_jacobian_w(site_id, dof_indices)¶
Compute world-frame Jacobians for one site and selected DoF columns.
- get_terrain_spawn_data()¶
Return backend-materialized terrain metadata on the cold path.
Backends without generated terrain support return
None. Callers should resolve this once during env initialization and cache the returned height-sampling callable for reset/reward hot paths.- Return type:
TerrainSpawnData|None
- render()¶
Render one frame through a backend-native interactive renderer.
- Raises:
RenderClosedError – If the user closed the render window.
- Return type:
- set_mocap_pose_tensor(body_name, env_indices, poses)¶
Write selected mocap poses through the declared tensor lifecycle.
- set_physics_state(state)¶
Restore a snapshot produced by
get_physics_state.Backends implementing this must refresh their host caches so state and sensor getters stay consistent with the restored physics state.
- set_pre_step_control(fn)¶
Register an env-owned policy-control to physics-control converter.
The callback receives
(backend, ctrl)so owner code can read the backend’s freshly-updated sensor contract before every physics substep. It must return either backend-native actuator control with the same shape, or aPreStepControlOutputcarrying that control plus an optional per-substep body wrench. A returned wrench is recomposed from scratch every substep and cleared when thestep()call finishes. Position-actuator envs leave this unset and keep the direct control path.
- backend_type: str¶