kai0-towel-advantage-base-stage1-mix

Kai0 Towel Advantage Estimator — Base + Stage 1 Mix

This repository contains a trained Advantage Estimator checkpoint from Kai0 for a dual-arm towel-folding task.

The model estimates task progress differences from pairs of multi-view observations. It is intended for research on stage-aware robotic manipulation, checkpoint evaluation, dataset labeling, and downstream Advantage-Weighted Behavior Cloning (AWBC).

Checkpoint

• Kai0 config: ADVANTAGE_TORCH_KAI0_TOWEL_FOLD
• Experiment: towel_base_stage1_mix
• Training step: 10000
• Initialization: PyTorch Pi0.5 base checkpoint
• Precision: bfloat16
• Global batch size: 64
• Seed: 42
• Value loss weight: 1.0
• Action loss weight: 0.0

The training sampler first selected one of two dataset groups with equal probability:

• Group A: towel-folding base
• Group B: towel-folding stage1

Sampling within each selected group was uniform over its frames. The model used front and left-wrist camera views; robot state was masked for this configuration.

The training prompt was:

fold the towel on the table and put it into the basket

Files

• model.safetensors: trained Advantage Estimator weights.
• metadata.pt: Kai0 training configuration and checkpoint metadata (global_step=10000).

The optimizer state is intentionally not included, so this repository provides inference/evaluation weights rather than a complete training-resume checkpoint.

Usage

Clone the Kai0 repository and install its pinned environment, then download this repository and pass its local path as the Advantage Estimator checkpoint directory. For example:

uv run python scripts/visualize_advantage_checkpoint.py \
  --ckpt-dir /path/to/kai0-towel-advantage-base-stage1-mix \
  --config-name ADVANTAGE_TORCH_KAI0_TOWEL_FOLD \
  --dataset /path/to/lerobot/dataset \
  --output-dir /path/to/output

This is a Kai0-specific checkpoint and is not a standalone Hugging Face Transformers model. Loading and preprocessing must follow the Kai0 Advantage Estimator implementation and the matching training configuration.

Intended use and limitations

This checkpoint is intended for non-commercial robotics research. Its predictions depend on the observation layout, camera viewpoints, task prompt, data distribution, and preprocessing used during training. Validate it on held-out data before using its scores for labeling, OOD/stall heuristics, or policy training. It should not be treated as a general robotic reward model or as a safety guarantee.

License

This checkpoint is released under the CC BY-NC-SA 4.0 license.