Convenience method for initializing all parameters, necessary due to the quirks of linen.

Update the running statistics with a batch of vectors. Args: vectors (np.ndarray): An array where all dimensions except

Helper method to apply gradient checkpointing if enabled.

Create a model with the given parameters.

Log detailed memory usage information.

Flatten a nested dictionary. Uses '/' as the separator.

Compute and return the statistics of the vectors processed so far. Returns: dict: A dictionary containing the computed s

Loads the model weights. Args: params: Parameters of the model. This is a nested structure of array-like objects that

Pad an array to the target dimension with zeros along the specified axis.

Append transforms to the group and return a new group. Args: inputs: Appended to the *end* of the current input transforms.

Get the mtime of a given date at midnight UTC.

Used to convert between aloha and pi joint angles.

Helper method to prepare 4D attention masks for transformer.

Set folder permission to be read, write and searchable.

chmod requires executable permission to be set, so we skip if the permission is already match with the target.

Get the data config for this data loader.

Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors)

Returns the freeze filter based on the model config.

Copied from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_a

Adapted from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_

Record a timing measurement for the given key.

Reset the agent to its initial state.

Applies RoPE positions [B, L] to x [B, L, H, D].

Applies RoPE positions [B, L] to x [B, L, H, D].

Download a file from a remote filesystem to the local cache, and return the local path.

Since we are sharing cache directory with containerized runtime as well as training script, we need to ensure that the cache directory has the cor

Applies gated residual connection with optional gate parameter. Args: x: Input tensor (residual) y: Output tensor to be

Get bins per dimension based on codebook size, from the original FSQ paper.

Merges the loaded parameters with the reference parameters. Args: loaded_params: The parameters to merge. params: The reference p

Helper method to preprocess observation.

Embed images with SigLIP and language tokens with embedding layer to prepare for PaliGemma transformer processing.

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

This method defines the mapping between unstructured data (i.e., nested dict) to the structured Observation format.

Obtains image last hidden states from the vision tower and apply multimodal projection. Args: pixel_values (`torch.Float

Returns the input specification for the model. Values are jax.ShapeDtypeStruct.

Adapted from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, he

Rotates half the hidden dims of the input.

Runs the runtime loop continuously until stop() is called or the environment is done.

Convert the Observation to a nested dict.