Unflatten a flattened dictionary. Assumes that '/' was used as a separator.

Adjust histograms when min or max changes.

Compute quantiles based on histograms.

Download a file or directory from GCS using gsutil if available, otherwise fall back to gcsfs.

Get bins per dimension according to the Lookup-Free Quantization paper (2 bins per dimension)

Initializes and returns the LIBERO environment, along with the task description.

Initialize KV cache

Loads and validates the weights. Returns a loaded subset of the weights.

Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_util

Set to position mode and do position resets: first close then open. Then change back to PWM mode NOTE: This diverges from the original Aloha

Replicates tf.image.resize_with_pad for one image using PIL. Resizes an image to a target height and width without distortion by padding with zero

Runs a single episode.

Set all files to be read & writable, if it is a script, keep it as a script.

Invalidate the cache if it is expired. Return True if the cache was invalidated.

A single step of the runtime loop.

Update KV cache with new values

Update histograms with new vectors.

Take an action in the environment.

Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): T

Convert PI0 JAX checkpoint to PyTorch format. Args: checkpoint_dir: Path to the JAX checkpoint precision: Model precision (f

Create a default policy for the given environment.

Create a policy from the given arguments.

Create an RLDS data loader for training. Note: This data loader requires some extra dependencies -- see examples/droid/README_train.md Args:

Computes sine-cosine positional embedding vectors for scalar positions.

Create a data loader for training. Args: data_config: The data configuration. action_horizon: The action horizon. batch_s

Create a dataset for training.

Converts a string like "B" or "B/32" into a params dict.

Apply one denoising step of the noise `x_t` at a given timestep.

Deserialize the running statistics from a JSON string.

Query the agent for the next action.

r""" Returns: image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by

Get the latest checkpoint step number from a checkpoint directory.

Query the environment for the current state.

Get a safe dtype for the given device type.

Get statistics for the given key.

Enable gradient checkpointing for memory optimization.

Infer actions from observations.

Custom logging format for better readability.

Initialize wandb logging.

This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. This m

Allow the environment to signal that the episode is complete. This will be called after each step. It should return `True` if the episode is

Converts input from left-align to right-aligned.

Load the latest checkpoint and return the global step.