Create a data loader for training. Args: config: The training configuration. sharding: The sharding to use for the data loader. I

Create a policy from a trained checkpoint. Args: train_config: The training config to use to create the model. checkpoint_dir: Th

Apply FSDP sharding to a pytree of arrays based on the mesh shape. Args: pytree: A pytree to be apply sharding specified by the mesh, not

Get a config by name.

Returns config for specified gemma variant.

Returns config for specified gemma variant.

Returns the freeze filter based on the model config.

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

Load the normalization stats from a directory.

Creates a random input example for the Aloha policy.

Make a boolean mask for the given dimensions. Example: make_bool_mask(2, -2, 2) == (True, True, False, False, True, True) make_bo

Creates a random input example for the Droid policy.

Creates a random input example for the Libero policy.

Download a file or directory from a remote filesystem to the local cache, and return the local path. If the local file already exists, it will be

The model type.

A higher-order function to JIT-compile `nnx.Module` methods, freezing the module's state in the process. Why not `nnx.jit`? For some reason, naiv

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

Preprocess the observations by performing image augmentations (if train=True), resizing (if necessary), and filling in a default image mask (if ne

Reset the policy to its initial state.

Reset the environment to its initial state. This will be called once before starting each episode.

Replicates tf.image.resize_with_pad for multiple images using PIL. Resizes a batch of images to a target height. Args: images: A batch of

Replicates tf.image.resize_with_pad. Resizes an image to a target height and width without distortion by padding with black. If the image is float

Restores unstructured params PyTree from a checkpoint. This works with checkpoints saved with `save_state` during openpi training (see `training/

Runs the runtime loop in a new thread.

Save the normalization stats to a directory.

Plumbing the mesh deep into the module tree is extremeley cumbersome; until the JAX team lands a better API, a custom context manager like this on