Create a dataset pipeline applying configured transforms.

Convert rotations given as angle-axis to quaternions. Args: angle: The angle turned anti-clockwise in radians around the vector's directi

Resolve a checkpoint path to a concrete checkpoint location. Supports checkpoint directories (checkpoint_*/model.pt) and hf: paths. Returns:

Multiply two quaternions together. Args: q1: The first quaternion in (w, x, y, z). Shape is (..., 4). q2: The second quaternion i

Instantiate transform specs into concrete transform objects.

Load a checkpoint into `model`. `path` is a checkpoint directory containing `model.pt`, or a direct path to `model.pt`. Returns (missing_keys

Computes the conjugate of a quaternion. Args: q: The quaternion orientation in (w, x, y, z). Shape is (..., 4). Returns: The

Return total bytes and float-dtype bytes for a list of tensors.

Convert rotations given as quaternions to axis/angle. Args: quat: The quaternion orientation in (w, x, y, z). Shape is (..., 4).

Compute training loss between predicted and target actions.

Manual autoregressive decoding loop with explicit KV cache updates.

Helper to decode a single token tensor.

Import and return a class/function given a fully-qualified path string.

Convert rotations given as quaternions to rotation matrices. Args: quaternions: The quaternion orientation in (w, x, y, z). Shape is (...

Normalizes a given input tensor to unit length. Args: x: Input tensor of shape (N, dims). eps: A small value to avoid division by

Project a rotation matrix to the 6D representation (first two rows). Args: rotation: tensor of shape (..., 3, 3) Returns: ten

Convert env RGB image (H, W, 3) to CHW float32 in [0, 1].

Send a response to a specific client via ROUTER socket.

Manual decoding loop mirroring test_qwen3vl_generation.decode.

Construct SuffixInput from VLMOutputs for caching purposes.

Do a full inference forward and compute the action (batch_size x num_steps x num_motors)

Merge saved cfg.yaml from a checkpoint run directory into `cfg`. Only `policy` and `data` groups are merged to keep runtime overrides intact.

Apply one denoising step of `noisy_actions` at a given timestep.

Prefill pass that returns cache, cache_position, attention_mask, and next token.

Apply an inverse quaternion rotation to a vector. Args: quat: The quaternion in (w, x, y, z). Shape is (..., 4). vec: The vector

Repeat a tensor along a new leading dimension then flatten.

Resolve placeholders like '{data.attr}' or '{policy.attr}' in a string/Path. Unknown placeholders are left untouched; only 'data.*' and 'policy.*

Sample actions in evaluation/serving mode.

Send a response back to the client. The response is queued and will be sent by the background thread. The client_id is extracted from

De-normalizes a given input tensor from range of [-1, 1] to (lower, upper). .. note:: It uses pytorch broadcasting functionality to deal

Recursively return total bytes and float-dtype bytes for optimizer state.

Return the rotation matrices for one of the rotations about an axis of which Euler angles describe, for each value of the angle given. Args:

Build a (episode_index, frame_index) -> byte offset index for jsonl bboxes. Only stores offsets (and keys) to avoid loading bbox payloads int

Decode raw multipart request frames into a dict.

Embed state, noisy_actions, timestep to prepare for Expert Gemma processing.

Encode a response payload into multipart frames.

Split values can be either lists of episode ids or range strings like \"0:1376\".

Split values can be either lists of episode ids or range strings like \"0:1376\".