unisim.backend.mujoco.tensor_io

Persistent packed tensor transfers for the MuJoCo/MJBatch host bridge.

Classes

MuJoCoHostBridgeTransferPlan

One stable layout for MuJoCo's unavoidable CPU tensor boundaries.
class unisim.backend.mujoco.tensor_io.MuJoCoHostBridgeTransferPlan[source]

Bases: HostBridgeTransferPlan

One stable layout for MuJoCo’s unavoidable CPU tensor boundaries.

CPU physics remains authoritative. The plan removes per-field allocations, metadata lookup, and per-sensor device transfers from the hot path. Each public operation performs at most one packed transfer on the current Torch stream and synchronizes only that stream when crossing to or from CUDA.

Parameters:

  • backend (MuJoCoBackend)

  • spec (TensorIOSpec)

__init__(backend, spec)[source]
Parameters:

  • backend (MuJoCoBackend)

  • spec (TensorIOSpec)

property spec: TensorIOSpec

Return the immutable layout request used to compile this plan.

property transfer_stats: dict[str, int]

Return cumulative semantic transfer and synchronization counters.

write_control(ctrl)[source]

Stage one complete control tensor on the host.

Parameters:

ctrl (Any)

Return type:

None

step(nsteps=1)[source]

Run CPU physics with the tensor control staged by write_control.

Parameters:

nsteps (int)

Return type:

dict | None

read_state_sensors()[source]

Return persistent tensor views after one packed H2D read.

Return type:

Mapping[str, Any]

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

Pack selected reset rows once, D2H them, and commit CPU physics.

randomization is the torch-first TensorResetRandomizationPayload; its plane check and host NumPy conversion happen once in _host_set_state.

Parameters:
Return type:

dict | None

read_selected_state_sensors()[source]

Return full views after one packed selected-row post-reset H2D.

Return type:

Mapping[str, Any]

close()[source]

Release transfer staging ownership without closing CPU physics.

Return type:

None