Daniel-F/tempo-mot-uvt

TEMPO-MOT + UVT

Four fine-tuned pi0.5 policies, the data they were trained on, and the code to load, train and serve them. Everything resolves inside this folder.

checkpoints/<task>/model.safetensors   policy weights (7.1 GB each)
checkpoints/<task>/assets/             norm stats, read when serving
actvae/<task>/actvae.pt                frozen action-chunk VAE
assets/pi05_yam_tempo_mot_uvt_<task>/  norm stats, read when training
data/<repo_id>/                        the LeRobot dataset for each task
sam2_cache/<task>/                     precomputed SAM2 cue, one npz per episode
openpi/                                training and inference code
tools/precompute_sam2_tokens.py        build a SAM2 cache for a new dataset
tools/actvae/                          train the action-chunk VAE for a new task
scripts/                               finetune.sh, finetune.sbatch, check_checkpoint.py, eval_loss.py

task	config	episodes	action dim	cameras
dynamic_handover	pi05_yam_tempo_mot_uvt_dynamic_handover	90	14	head, left_wrist, right_wrist
dynamic_pour	pi05_yam_tempo_mot_uvt_dynamic_pour	23	7	head, right_wrist
spartan_balls	pi05_yam_tempo_mot_uvt_spartan_balls	54	7	head, right_wrist
ball_drop	pi05_yam_tempo_mot_uvt_ball_drop	31	7	head
push_ball_blk	pi05_yam_tempo_mot_uvt_push_ball_blk	50	7	head, right_wrist

Setup

Needs Python 3.11+, an NVIDIA GPU, and uv.

huggingface-cli download Daniel-F/tempo-mot-uvt --local-dir tempo-mot-uvt
cd tempo-mot-uvt/openpi
uv sync
# required: the temporal attention lives in a patched SigLIP and the model will
# not build without it
cp -rf src/openpi/models_pytorch/transformers_replace/* \
       .venv/lib/python3.11/site-packages/transformers/

Re-run that copy after editing anything under transformers_replace/.

export TEMPO_UVT_VENV=$PWD/.venv    # the only variable you have to set

TEMPO_UVT_ROOT, HF_LEROBOT_HOME and TEMPO_SAM2_CACHE_ROOT default to this folder. Set them only to train against data or caches kept elsewhere.

Load

python scripts/check_checkpoint.py ball_drop

Builds the model from the task's config, loads its weights, and samples one action chunk. Expect missing=0 unexpected=0 and a finite (1, 16, action_dim) chunk.

In process:

import openpi.training.config as config
from openpi.policies import policy_config

cfg = config.get_config("pi05_yam_tempo_mot_uvt_spartan_balls")
policy = policy_config.create_trained_policy(cfg, "checkpoints/spartan_balls")

action_chunk = policy.infer({
    "state": state,                        # (A,) float32
    "images": {"head": head_rgb, "right_wrist": right_rgb},   # HWC uint8
    "prompt": "spartan balls",
    "sam2_tokens": sam2_tokens,            # (64, 256) float32
})["actions"]                              # (16, A) raw units

sam2_tokens is SAM2.1-tiny's memory-attention output for the head frame, (256, 8, 8) read as 64 tokens x 256. The backbone takes three camera slots; slots a dataset does not have are zero-filled and masked off.

Or serve it over a socket:

cd openpi
python scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_yam_tempo_mot_uvt_spartan_balls \
    --policy.dir=../checkpoints/spartan_balls --port=8000

from openpi_client import websocket_client_policy

client = websocket_client_policy.WebsocketClientPolicy(host="localhost", port=8000)
action_chunk = client.infer({...})["actions"]

Train

bash scripts/finetune.sh spartan_balls

Continues from the checkpoint of the same task. Checkpoints go to runs/<task>/<config>/<exp_name>/<step>/; re-running the same command resumes from the newest one.

Knobs: EXP_NAME, BATCH_SIZE (global), NPROC, CHECKPOINT_BASE_DIR. Extra flags pass straight through to the trainer:

BATCH_SIZE=64 bash scripts/finetune.sh ball_drop --num-train-steps=20000

Under SLURM, submit from the root of this folder:

sbatch --nodelist=<node> scripts/finetune.sbatch spartan_balls

To check a checkpoint's loss on its data without training:

python scripts/eval_loss.py spartan_balls --batches 25

A new task

Add a config next to the existing ones in openpi/src/openpi/training/config.py. It needs a dataset under data/, a SAM2 cache, norm stats, and an action-chunk VAE, in that order:

export SAM2_REPO=/path/to/sam2 SAM2_CKPT=/path/to/sam2.1_hiera_tiny.pt
python tools/precompute_sam2_tokens.py \
    --root data/<repo_id> --cache sam2_cache/<task>

cd openpi && JAX_PLATFORMS=cpu python scripts/compute_norm_stats.py \
    --config-name <config> && cd ..

python -m tools.actvae.train \
    --root data/<repo_id> \
    --norm-stats assets/<config>/<repo_id>/norm_stats.json \
    --out actvae/<task> \
    --n-tokens 2 --token-dim 32 --chunk 16 \
    --betas 1e-5 1e-4 1e-3 1e-2 --epochs 300

Use --n-tokens 2 for a 7-D arm and 4 for 14-D, so the latent stays a compression of the 16-step chunk rather than an expansion.

License

Apache 2.0, following openpi. These weights are fine-tuned from pi0.5, which contains Gemma; see openpi/LICENSE_GEMMA.txt for the terms that carry over.