OpenPI: JAX policy training

Train or fine-tune an OpenPI policy with the JAX training pipeline and an explicit upstream training/data configuration. Inspect the pinned source and setup before use.

openpi-scripts-train-py | Record version: source-22f98a06aaa6e68e0f725d727ab95398238d1cae8a46bb94e829f330758a9baf

Record JSON | Markdown | Pinned upstream source

Access

Fetch the pinned repository, verify this source checksum, and use upstream installation and invocation. This record does not install or execute code.

Hub preparation adapter: false. Follow the upstream setup and invocation contract.

Source and license

Repository: https://github.com/Physical-Intelligence/openpi
Revision: 215abfb217dbac7d5f1273282331b9b1866c0479
Path: scripts/train.py
Source SHA256: 22f98a06aaa6e68e0f725d727ab95398238d1cae8a46bb94e829f330758a9baf
Original authors: Physical Intelligence team
License: Apache-2.0
License scope: Selected first-party workflow Python source only; preserve all upstream copyright, author and per-file notices. This does not relicense model weights, datasets, Gemma or other third-party model code, robot firmware, simulator binaries, meshes, dependencies or external services; their separate terms and access requirements remain applicable.
License SHA256: 43070e2d4e532684de521b885f385d0841030efa2b1a20bafb76133a5e1379c1
Hub commit: 1d6d9eb134e397d910265a8d6e7d8f257f9e8d7f

Artifact

{
  "kind": "workflow",
  "format": "python",
  "workflow_role": "policy-training"
}

Requirements

{
  "python": "Python >=3.11 as declared by the pinned pyproject.toml",
  "packages": [
    "Pinned OpenPI checkout and openpi-client",
    "Upstream uv environment; JAX/CUDA or PyTorch backend as selected",
    "Environment-specific ALOHA, DROID or LIBERO dependencies when selected"
  ],
  "display": "Upstream supports Ubuntu 22.04. Policy inference/training requires suitable NVIDIA GPU memory and separately obtained checkpoints/data; real robot clients additionally require the exact upstream robot setup and a configured policy server.",
  "execution_domain": "computation",
  "source_imports": [
    "dataclasses",
    "etils.epath",
    "flax.nnx",
    "flax.training",
    "flax.traverse_util",
    "functools",
    "jax",
    "jax.experimental",
    "jax.numpy",
    "logging",
    "numpy",
    "openpi.models.model",
    "openpi.shared.array_typing",
    "openpi.shared.nnx_utils",
    "openpi.training.checkpoints",
    "openpi.training.config",
    "openpi.training.data_loader",
    "openpi.training.optimizer",
    "openpi.training.sharding",
    "openpi.training.utils",
    "openpi.training.weight_loaders",
    "optax",
    "platform",
    "tqdm_loggable.auto",
    "typing",
    "wandb"
  ]
}

Invocation

{
  "entrypoint": "uv run scripts/train.py",
  "notes": "Use the pinned upstream environment and CLI configuration, dataset/checkpoint paths and platform-specific README. Some upstream workflows can download assets or publish datasets; inspect flags before use. No checkpoint licensing, robot compatibility or runtime success is established by this source reference."
}

Explicit compatibility

[]

Missing environment fields remain unknown. Discovery labels do not establish robot compatibility.

Evidence

Level: source_reviewed; runtime verification: not_run.

Verified exact source, applicable license and author/setup document bytes at a full Git commit; parsed the Python AST and reviewed this guarded operational entrypoint and its explicit license notices. No upstream code, installation, robot command, dataset, mesh or checkpoint was executed or loaded.

Source review and contributor reports do not establish task completion. Preserve upstream attribution and licenses.