GPU & JAX Sims

GPU and JAX simulators (like Isaac Gym, Brax, or MJX) keep their state in device tensors that do not have stable host memory addresses.

Because zero-copy pointer reflection requires a stable CPU pointer, OmniLoop provides the SnapshotReflector. It reduces device tensors to host scalar summaries (mean, min, max) for telemetry without crashing your pipeline or forcing expensive device-to-host syncs where not needed.

Usage

from omniloop.integrations.gpu_snapshot import SnapshotReflector

reflector = SnapshotReflector(tracker)

reflector.snapshot({
    "joint_positions": sim.get_dof_position(),  # GPU tensor
    "joint_velocities": sim.get_dof_velocity(),
})

Note

This requires no extra pip packages. It works seamlessly with whatever tensor library you already have installed (PyTorch, JAX, or NumPy) via duck-typing.

When to use Snapshots vs Backends

• Use SnapshotReflector (this page) for device tensors and GPU memory where pointers change or don’t exist on the host.
• Use BackendAdapter (see MuJoCo Backend) for CPU contiguous f64 state where zero-copy pointer reflection is possible.