Phy-lab/aic_cable_insert_sft_openpi_pi05_pretrained_v3_bs128_mb16_step11040

AIC Cable Insert SFT OpenPI Pi0.5 UR5e

이 저장소는 AIC cable insertion 작업을 위해 UR5e 기준으로 학습한 RLinf/OpenPI Pi0.5 SFT checkpoint를 제공합니다.

학습 및 inference에 사용하는 RLinf OpenPI config 이름은 다음과 같습니다.

pi05_aic_cable_insert_ur5e

필요한 RLinf PR

이 checkpoint를 실행하려면 RLinf에 AIC UR5e OpenPI config, policy transform, normalization stats, inference script가 필요합니다.

해당 최소 변경은 아래 PR에 있습니다.

https://github.com/Phy-Lab-aic/RLinf/pull/6

PR이 merge되기 전에는 PR branch를 checkout해서 실행합니다.

git clone https://github.com/Phy-Lab-aic/RLinf.git
cd RLinf
git fetch origin pull/6/head:aic-ur5e-openpi-inference
git checkout aic-ur5e-openpi-inference

RLinf 환경 준비

RLinf venv 기준으로 실행합니다.

python -m venv .venv
source .venv/bin/activate
pip install -U pip
pip install -e ".[embodied,openpi]"

이미 RLinf venv가 준비되어 있다면 source .venv/bin/activate만 수행하면 됩니다.

Checkpoint 구성

각 checkpoint는 별도 폴더로 저장되어 있습니다.

global_step_2208
global_step_4416
global_step_6624
global_step_8832
global_step_11040

기본 inference checkpoint는 global_step_11040입니다.

각 checkpoint 폴더에는 다음 파일들이 포함되어 있습니다.

actor/model_state_dict/full_weights.pt
actor/dcp_checkpoint/
Phy-lab/pretrained_dataset_v3_lerobot_v21_compat/norm_stats.json

norm_stats.json은 OpenPI 입력/출력 normalization에 필요합니다.

입력 Observation 형식

AIC UR5e policy는 RGB 카메라 3개와 7D state를 입력으로 받습니다.

env_obs = {
    "main_images": center_rgb_image,
    "wrist_images": left_or_wrist_rgb_image,
    "extra_view_images": right_or_extra_rgb_image,
    "states": ur5e_state_tensor,
    "task_descriptions": ["Insert the SFP-to-SC cable into the target port"],
}

7D state/action 순서는 다음과 같습니다.

[shoulder_pan, shoulder_lift, elbow, wrist_1, wrist_2, wrist_3, gripper]

policy 출력 action chunk의 shape은 다음과 같습니다.

[batch, 5, 7]

마지막 차원의 7번째 값은 gripper command입니다.

바로 실행

RLinf PR branch 안에서 아래 명령을 실행하면 Hugging Face에서 global_step_11040 checkpoint를 내려받고, zero image/zero state로 inference 경로를 확인합니다.

python examples/sft/aic_ur5e_openpi_inference.py \
  --model-path Phy-lab/aic_cable_insert_sft_openpi_pi05_pretrained_v3_bs128_mb16_step11040 \
  --checkpoint global_step_11040

다른 checkpoint를 사용하려면 --checkpoint 값을 바꿉니다.

python examples/sft/aic_ur5e_openpi_inference.py \
  --model-path Phy-lab/aic_cable_insert_sft_openpi_pi05_pretrained_v3_bs128_mb16_step11040 \
  --checkpoint global_step_8832

실제 UR5e 입력으로 실행

현재 UR5e 7D state를 JSON 파일로 준비합니다.

[-0.49, -1.74, -1.55, -1.56, 1.88, 1.11, 0.007]

그 다음 RGB 이미지 3장과 state JSON을 지정해서 실행합니다.

python examples/sft/aic_ur5e_openpi_inference.py \
  --model-path Phy-lab/aic_cable_insert_sft_openpi_pi05_pretrained_v3_bs128_mb16_step11040 \
  --checkpoint global_step_11040 \
  --center-image /path/to/center.png \
  --left-image /path/to/left_or_wrist.png \
  --right-image /path/to/right_or_extra.png \
  --state-json /path/to/ur5e_state.json

출력은 JSON 형식입니다.

{"actions": [[[... seven action values ...]]]}

PR 범위

PR #6은 Hugging Face checkpoint를 UR5e OpenPI inference로 실행하기 위한 최소 변경만 포함합니다. IsaacLab, Gazebo, RL runtime integration은 후속 PR에서 다룹니다.