unisim.backend.mjwarp.randomization

Per-world model-field expansion for the independent mjwarp backend.

MuJoCo Warp declares the leading dimension of randomizable Model fields as "*": every kernel reads field[worldid % field.shape[0]], so tiling a field from (1, ...) to (nworld, ...) yields per-world semantics without kernel changes. Expansion replaces the array allocation, so it is a cold-path operation: the backend expands the declared set once during construction, before CUDA graph capture, and all later DR writes are in-place assign uploads into the same fixed-address arrays (graph-safe).

The field list is ported from mjlab’s expand_model_fields usage and checked against the pinned mujoco-warp 3.11 Model dataclass. Derived fields are expanded alongside their inputs because the set_const* recompute kernels size their launch grid from the output field’s leading dimension.

Functions

expand_model_fields(warp, model, nworld)

Tile the declared DR model fields from (1, ...) to (nworld, ...).

Classes

PrimitiveGeomBounds

Cold-path geometry classification for batch-only bounds recomputation.
class unisim.backend.mjwarp.randomization.PrimitiveGeomBounds[source]

Bases: object

Cold-path geometry classification for batch-only bounds recomputation.

Uses official MuJoCo primitive geometry formulas. Meshes, planes, height fields and SDFs may appear in dense payloads unchanged, but resizing them is unsupported. No model/XML metadata is consulted during reset.

Parameters:
__init__(geom_types, geom_enum)[source]
Parameters:
compute(sizes, previous_sizes, previous_rbound, previous_aabb)[source]

Validate a dense size update and derive detached selected-row bounds.

Parameters:
Return type:

tuple[ndarray, ndarray]

unisim.backend.mjwarp.randomization.expand_model_fields(warp, model, nworld)[source]

Tile the declared DR model fields from (1, ...) to (nworld, ...).

Returns the names actually expanded. Fields already per-world (e.g. from a prior expansion) are skipped, mirroring mjlab’s guard. A single-world backend keeps the shared arrays untouched.

Parameters:
Return type:

tuple[str, ...]