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

Config for training on DROID, using RLDS data format (for efficient training on larger datasets).

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

Load a policy from a trained checkpoint.

Creates model transforms for standard pi0 models.

Transformer MLP / feed-forward block.

Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`SiglipEncoderLayer`]. Args:

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).

Gated Linear Unit with GELU (GeGLU) activation function. GeGLU is a Flax layer that combines a linear transformation with a GELU activation fu

Multi-headed attention from 'Attention Is All You Need' paper

A group of transforms.

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

Example data config for custom DROID dataset in LeRobot format. To convert your custom DROID dataset (<10s of hours) to LeRobot format, see e

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

r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token pr

r""" past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):

Shallow wrapper around the DROID data loader to make it compatible with openpi. All batching already happens in the DROID dataset, so we don't ne

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

Multi-headed attention from 'Attention Is All You Need' paper

Multihead Attention Pooling.

r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): Contrastive loss for image-text simila

Records timing measurements for different keys.

Torch data loader implementation.

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.

Command line arguments.

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

Standard RT-2 / OpenVLA style binning tokenizer.

Transformer block.

Transformer block.