OakEmu-RL v0.1.0

Appearance

ObservationTransferWorkspace ​

python
ObservationTransferWorkspace(device: torch.device, batch_shape_hint: tuple[int, ...] | None = None) -> None

Reuse a device-side observation buffer for repeated inference uploads.

CUDA inference would otherwise allocate a fresh device tensor for every numpy -> torch -> cuda transfer. This workspace keeps a reusable observation buffer and copies the current batch into a slice of it, staging host data through a matching pinned buffer so the upload runs as a fast non-blocking DMA copy instead of a pageable transfer.

Initialize the observation transfer workspace.

Parameters:

  • device (device) – Target device for observation tensors.
  • batch_shape_hint (tuple[int, ...] | None) – Optional maximum batch shape to pre-size the reusable buffer.

Source

Methods ​

to_device ​

python
to_device(observations: torch.Tensor | np.ndarray) -> torch.Tensor

Move observations to the target device, reusing a CUDA buffer when possible.

Parameters:

  • observations (Tensor | ndarray) – Observations as a torch tensor or numpy array.

Returns:

  • Tensor – A tensor on the target device.

Source