WoWjx/so101-stack2-pi05-full-relative

SO101 Stack2 π0.5 — Full FT / Relative Actions

This revision contains step 4228 inference artifacts. Candidates for real-robot evaluation; physical Stack2 and Stack3 success has not been established.

Model and robot

SO101; fine-tuned from lerobot/pi05_base, base revision b211f3d44c36b6acfcf7ae94a64e8e96f75a64ba. Two RGB cameras: observation.images.front and observation.images.wrist, original 640×480, resize-with-pad to 224×224. State 6, action 6: first five arm joints relative to current raw state, gripper absolute. Paired official processors unnormalize and restore absolute arm commands while preserving the absolute gripper. Use LeRobot commit ff71cae1ae2d09fd035553c35da65888ed6c8304.

The PaliGemma/Gemma components inherit the base model's Gemma terms: https://ai.google.dev/gemma/terms . No additional license grant is claimed.

Dataset

WoWjx/so101-stacking-corpus, publication snapshot 0207d53312ec111a1355436688ee83360c927818. Stack2 only. Train: 101 episodes / 58,453 frames; validation: 14 / 8,473; held-out Stack2: 21 / 11,342. Collection batches isolated; six directed red/blue/yellow relations in all splits. Stack3's 64 episodes and all temporal sidecar annotations were excluded from training, normalization statistics and validation. Normalization fitted on train only. Task text: stack the <source> block on the <target> block.

Training

Full fine-tuning; seed 1000; batch 64; BF16; gradient checkpointing; LR 2.5e-5; AdamW betas (0.9,0.95), eps 1e-8, weight decay 0.01, gradient clipping 1.0. Cosine decay over 30,000 updates, warmup 1,000, terminal LR 2.5e-6. Chunk size 50; n_action_steps 10. Validation every 500; checkpoint every 2,500 plus graceful final save. Requested max_steps 30,000; actual stop 4,228 due to the ten-hour campaign deadline. No image/prompt augmentation.

A100-SXM4-80GB; CUDA 12.6; PyTorch 2.7.1+cu126; Python 3.13.9; Transformers 5.5.4; LeRobot 0.6.2 at the pinned commit above. Approximately 7.3s/update, 8h45m active training; steady GPU utilization 99–100%; peak allocated VRAM 60.56 GiB.

Available revisions

• main: step 4228.
• step-2500: step 2500.

Inference artifacts only. Complete local optimizer/RNG/scheduler state exists for step 4228 and is not included publicly. Never use the public bundle as if it were a full optimizer resume checkpoint.

Offline metrics

Final step-4228 train loss 0.02439656; last-25 logged mean 0.015489. Best recorded validation flow-matching loss 0.0537 at unsaved step 1500. Step 2500 validation 0.0614. Latest preceding validation step 4000: 0.0704. These stochastic offline losses are not robot success rates, and do not establish which checkpoint is physically better.

Intended evaluation and limitations

Stack2 held-out physical evaluation; Stack3 B0 direct long instruction; Stack3 B1 oracle human language handoff, without retraining or Stack3 demonstrations. B1 supplies privileged stage information and is not autonomous planning. Only one seed trained; absolute-action and expert-only ablations were not trained. No relative-vs-absolute, Full-FT-vs-expert-only, seed, Stack2 success, Stack3 success or composition claim is made.

Code / reproducibility

GitHub source publication pending: the prepared local publication commit is 24d1b56d7ed4c1dbfad86baaf391cc7e9e053124 on pi05-stack2-training; it is not yet available on GitHub because repository write authentication is incomplete. Existing GitHub main does not yet contain this experiment. Model artifacts and the standalone load below are already public and verified.

Project: MmMaiIIi/SO101. Exact publication commit will be pinned after GitHub publication; recorded pre-training SO101 commit: f76f8cdc43a4311ae73a78cf74b9c80ab7701ee9. Prepared source files (not yet pushed) are in training/pi05_stack2/; report in reports/; B0/B1 protocol in eval/stack3/.

Download a specific revision using huggingface_hub.snapshot_download, then use the standalone load below, or the project's strict load_model.py loader after source publication. The pinned upstream PI05 loader does not reliably forward the named revision into its weight download; snapshot-first loading avoids mixing branch weights. The saved processor tokenizer is a relative artifact, not a cluster path. The config tokenizer path is resolved against the downloaded snapshot by the loader.

Standalone inference-bundle load

Use the pinned LeRobot commit and training package versions above. This snapshot-first code does not execute hardware. Set revision="step-2500" to load that candidate.

from pathlib import Path
from huggingface_hub import snapshot_download
from lerobot.configs import PreTrainedConfig
from lerobot.policies.pi05.modeling_pi05 import PI05Policy
import lerobot.policies.pi05.processor_pi05  # Register saved processor steps.
from lerobot.processor.factory import load_pretrained_policy_processors
from lerobot.processor.relative_action_processor import bind_relative_anchor

revision = "main"
path = Path(snapshot_download("WoWjx/so101-stack2-pi05-full-relative", revision=revision))
cfg = PreTrainedConfig.from_pretrained(path)
cfg.device = "cpu"
cfg.text_tokenizer_name = str(path / "tokenizer")
loaded = []
original = PI05Policy.load_state_dict
def strict_load(self, *args, **kwargs):
    result = original(self, *args, **kwargs)
    if result.missing_keys or result.unexpected_keys:
        raise RuntimeError(str(result))
    loaded.append(True)
    return result
PI05Policy.load_state_dict = strict_load
try:
    policy = PI05Policy.from_pretrained(path, config=cfg, local_files_only=True)
finally:
    PI05Policy.load_state_dict = original
if not loaded:
    raise RuntimeError("Pretrained weights were not loaded")
preprocessor, postprocessor = load_pretrained_policy_processors(
    path, preprocessor_overrides={"device_processor": {"device": "cpu"}}
)
bind_relative_anchor(policy, preprocessor)
policy.eval()

The paired saved postprocessor restores absolute arm commands and preserves the absolute gripper. The official anchor binding retains the raw-state reference throughout an action queue. No untracked custom action transform is required.

Verified public download

Both revisions were downloaded anonymously into a fresh cache, fully SHA256-checked, then strictly loaded on CPU (813 tensors) with paired-processor and arm/gripper round-trip checks. The cluster cannot directly reach the HF API, so verified public download files were transferred through SSH into independent directories. This is a model/processor integrity check, not physical robot evaluation.