Language Language English 简体中文

unisim.backend.mujoco.tensor_io.MuJoCoHostBridgeTransferPlan

class unisim.backend.mujoco.tensor_io.MuJoCoHostBridgeTransferPlan

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)

Methods

__init__(backend, spec)
apply_reset(env_indices, qpos, qvel[, ...])	Pack selected reset rows once, D2H them, and commit CPU physics.
close()	Release transfer staging ownership without closing CPU physics.
read_selected_state_sensors()	Return full views after one packed selected-row post-reset H2D.
read_state_sensors()	Return persistent tensor views after one packed H2D read.
step([nsteps])	Run CPU physics with the tensor control staged by write_control.
write_control(ctrl)	Stage one complete control tensor on the host.

Attributes

spec	Return the immutable layout request used to compile this plan.
transfer_stats	Return cumulative semantic transfer and synchronization counters.
last_timing

__init__(backend, spec)

Parameters:

• backend (MuJoCoBackend)

• spec (TensorIOSpec)

last_timing: dict[str, float]

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)

Stage one complete control tensor on the host.

Parameters:

ctrl (Any)

Return type:

None

step(nsteps=1)

Run CPU physics with the tensor control staged by write_control.

Parameters:

nsteps (int)

Return type:

dict | None

read_state_sensors()

Return persistent tensor views after one packed H2D read.

Return type:

Mapping[str, Any]

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

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:

• env_indices (Any)

• qpos (Any)

• qvel (Any)

• randomization (Any | None)

Return type:

dict | None

read_selected_state_sensors()

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

Return type:

Mapping[str, Any]

close()

Release transfer staging ownership without closing CPU physics.

Return type:

None