Determines the location of assets (e.g., norm stats) that will be used to set up the data pipeline. These assets will be replicated inside the ch

Load a policy from a trained checkpoint.

This config is used to configure transforms that are applied at various parts of the data pipeline. For your own dataset, you can copy this c

Creates model transforms for standard pi0 models.

Transformer MLP / feed-forward block.

Adds learned positional embeddings to the inputs. Attributes: posemb_init: positional embedding initializer.

Attention module.

Attention module.

A CheckpointHandler for calling an arbitrary function asynchronously. Only for saving, not for restoring.

A composite transform that applies a sequence of transforms in order.

Embedder module.

Embedder module.

Transformer Model Encoder for sequence to sequence translation.

Single transformer encoder block (MHSA + MLP).

A group of transforms.

Identity layer, convenient for giving a name to an array.

Multihead Attention Pooling.

Transformer MLP / feed-forward block.

Holds observations, i.e., inputs to the model. See `Observation.from_dict` to see the expected dictionary form. This is the format that shoul

Environment for real robot bi-manual manipulation Action space: [left_arm_qpos (6), # absolute joint position

ViT model.

Repacks delta actions into absolute action space.

Wraps a policy to return action chunks one-at-a-time. Assumes that the first dimension of all action fields is the chunk size. A new inferen

AdamW optimizer.

An Agent is the thing with agency, i.e. the entity that makes decisions. Agents receive observations about the state of the world, and return act

Inputs for the Aloha policy. Expected inputs: - images: dict[name, img] where img is [channel, height, width]. name must be in EXPECTED_CAMER

Outputs for the Aloha policy.

An environment for an Aloha robot on real hardware.

An environment for an Aloha robot in simulation.

Arguments for the serve_policy script.

Base class for all model implementations. Specific models should inherit from this class. They should call super().__init__() to initialize the sh

Configuration shared by all models. Specific models should inherit from this class, and implement the `create` method to create the corresponding

Transformer block.

Transformer block.

Loads an entire set of weights from a checkpoint. Compatible with: trained checkpoints: example: "./checkpoints/<config>/<exp>/<ste

Cosine decay schedule with warmup.

Interface for a data loader.

Interface for a dataset with random access.

Use the default policy for the given environment.

Repacks absolute actions into delta action space.

tokens:\n |<------------prefix(w=2048)----------->|<-----------suffix(w=1024)------------->|\n |<-images->|<-language->|<-effort(llm)->|<

Einsum with LoRA support. Can be used as a drop-in replacement for the Gemma Einsum.

Transformer Model Encoder for sequence to sequence translation. Attributes: num_layers: number of layers mlp_dim: dimension of the ml

Transformer encoder layer. Attributes: inputs: input data. mlp_dim: dimension of the mlp on top of attention block. dtype: the

Supported environments.

Supported environments.

An Environment represents the robot and the environment it inhabits. The primary contract of environments is that they can be queried for observa

Feed forward module.

Feed forward module.

This class is used to convert inputs to the model to the expected format. It is used for both training and inference. For your own dataset,

This class is used to convert outputs from the model back the the dataset specific format. It is used for inference only. For your own d

Configuration for LoRA.

Supported model types.

Transformer model, supporting a mixture of different weights for different tokens.

gemma model.

Loads weights from the official PaliGemma checkpoint. This will overwrite existing weights with similar names while keeping all extra weights int