realrobot-openpi-new-knife-v3 / inference_setup.md
haohw's picture
Upload inference_setup.md with huggingface_hub
4eab68e verified
|
Raw History Blame Contribute Delete
3.35 kB

Inference Setup for new_v3 Foundation + Edited Checkpoints

This repo's foundation ckpt (mixed_new_knife_v3/best_val/600/) and all edits/* ckpts share the same model arch and same norm_stats (real_pass_knife_new_mixed_train). They are not compatible with the original pi05_real_pk_mixed_v2 config's norm_stats (differences up to ~28 % on q99 — predictions will be visibly off if you use the wrong norm stats).

Setup (one-time, on the inference machine)

1. Add the TrainConfig entry to openpi

Open openpi/src/openpi/training/config.py and paste the contents of inference_config_snippet.py (in this folder) into the _CONFIGS list. The entry registers pi05_real_pk_mixed_new_v3 which points the data loader at the new norm_stats source.

2. Place norm_stats where openpi expects them

# Pull the bundled norm_stats from the foundation ckpt (already in this repo)
mkdir -p ~/openpi/assets/pi05_real_pk_mixed_new_v3
hf download haohw/realrobot-openpi-new-knife-v3 \
    --include "mixed_new_knife_v3/best_val/600/assets/real_pass_knife_new_mixed_train/norm_stats.json" \
    --local-dir /tmp/_pull
mv /tmp/_pull/mixed_new_knife_v3/best_val/600/assets/real_pass_knife_new_mixed_train \
   ~/openpi/assets/pi05_real_pk_mixed_new_v3/

Verify: ~/openpi/assets/pi05_real_pk_mixed_new_v3/real_pass_knife_new_mixed_train/norm_stats.json exists.

3. (Optional) Get the LeRobot eval dataset

If you want to run scripts/eval_* or load batches at inference, materialize the eval set:

hf download rollingoat/mutil_mode_data new_knife_lerobot_all.zip \
    --repo-type dataset --local-dir /mnt/data/new_knife_downloads
# extract + symlink as `~/.cache/huggingface/lerobot/real_pass_knife_mixed_v2_eval`
# (see the original build_pk_v2_split.py for the exact 10-ep split if needed)

For pure inference (no batched eval), step 3 is not required — you can feed observations directly via openpi's policy serving API.

Running inference

Foundation:

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

cfg = _config.get_config("pi05_real_pk_mixed_new_v3")
policy = policy_config.create_trained_policy(
    cfg, checkpoint_dir="/path/to/mixed_new_knife_v3/best_val/600",
)
# action = policy.infer(observation)

Edited ckpt (any of the edits/* folders):

policy = policy_config.create_trained_policy(
    cfg, checkpoint_dir="/path/to/edits/action_v4_keep_left_step100",
)

(The same TrainConfig works for the foundation and all edits — they only differ in the params/ weights, not in arch or norm stats.)

Sanity check before robot rollout

import jax, numpy as np
from openpi.training import config as _config
from openpi.training import data_loader as _data_loader

cfg = _config.get_config("pi05_real_pk_mixed_new_v3")
dc = cfg.data.create(cfg.assets_dirs, cfg.model)
print("asset_id:", dc.asset_id)
# should print: real_pass_knife_new_mixed_train
print("norm_stats keys:", list(dc.norm_stats))
# should include 'state' and 'actions'

If asset_id prints real_pass_knife_mixed_v2_train instead of ..._new_mixed_train, the asset_id override didn't take — re-check the snippet was pasted with assets=AssetsConfig(asset_id=...) on the outer RealRobotDataConfig (not inside base_config).