Pi0.5 Checkpoint 아키텍처 분석 + Cascade Adapter 설계

TL;DR

아키텍처: PaliGemma-3B + two-pathway action expert ≈ ~3.5B params. ViT(27L, hidden 1152) → LLM(18L, hidden 2048) + 병렬 1024-dim action expert pathway. orbax/JAX 포맷 → openpi 라이브러리로만 로드 가능.
Action spec: 32-dim 중 11-dim active (arm OSC_POSE 6 + gripper 1 + base velocity 3 + torso 1). dims 12–31은 cross-embodiment padding (std=1.0, 의미 없음).
Cascade adapter: MemerPi05Policy(BasePolicy) 골격 작성 완료. 3-venv IPC 설계 (memer torch / pi05 JAX / robocasa sim), reset()에서 prior 비디오 → MemER FIFO, get_action()에서 subtask → Pi0.5 action chunk.
비용 추정: 257 scenario × 11분 = 47h sequential. 8 GPU SLURM batch → ~6h, 또는 30 scenario subset으로 ~5h 시작.

~3.5B

Params

Active Dims

47h

Sequential Cost

~6h

8-GPU Batch

1 배경 / 목적

scene-mem-benchmark의 memory_success(approach fixture ∧ door open)를 실제 측정하려면 low-level VLA인 Pi0.5-RoboCasa가 필요하다. 260513 오전 probe에서 MemER (dusting ckpt)의 vocabulary lock-in 문제를 확인했고, 같은 날 오후 Pi0.5+MemER LoRA cascade 통합 plan을 확정했다.

문제: Pi0.5 ckpt가 orbax/JAX 포맷이라 transformers.AutoModel로 바로 로드 안 됨. MemER(torch 2.11)와 완전히 다른 프레임워크 → process 분리 + IPC 필수. action dim 레이아웃을 정확히 알아야 proprio state 매핑이 가능.

이 문서는 cascade adapter 작성 전에 Pi0.5의 내부 구조(파라미터, action/state spec, 로딩 API)를 확정하고, 다음 세션에서 즉시 코딩에 들어갈 수 있는 골격을 작성하기 위해 만들어졌다. 소스: _CHECKPOINT_METADATA, params/_METADATA, assets/norm_stats.json 직접 분석.

2 작업 내용

Pi0.5 Checkpoint 구조

대상: robocasa/robocasa365_checkpoints의 pi05_pretrain_human300/multitask_learning/75000/ (44 files, 260513 오전 49s 다운 완료).

pi05_pretrain_human300/multitask_learning/75000/ ├── _CHECKPOINT_METADATA # JSON: openpi handlers + orbax PyTreeCheckpointHandler ├── assets/ │ └── norm_stats.json # state/action 정규화 상수 (각 32-dim) ├── params/ # JAX/orbax sharded (OCDBT format) — 모델 가중치 │ ├── _METADATA # JAX pytree 구조 정의 │ ├── _sharding │ └── ocdbt.process_0/ # OCDBT manifest └── train_state/ # Optax optimizer state (inference에 불필요)

아키텍처 분석 — params/_METADATA 기반

PaliGemma (Google) 기반 + Pi-style action expert:

컴포넌트	세부
Image encoder (PaliGemma.img.Transformer)	ViT 27 layers, hidden 1152, MLP 4304, 16 heads × 72 dim. Patch: 14×14 → 1152. Out projection: 1152 → 2048.
LLM (PaliGemma.llm)	Gemma-style, 18 layers, hidden 2048, vocab 257152, 8 heads × 256 dim
Two-pathway parallel attention	Path 1: 2048 dim (image+text expert). Path 2: 1024 dim (action expert). 각각 별도 q_einsum/kv_einsum + mlp + pre_attention_norm
Action head (flow matching)	time_mlp_in/out (1024→1024), action_in_proj (32→1024), action_out_proj (1024→32)
총 파라미터	PaliGemma-3B + action expert ~500M ≈ ~3.5B (bf16 ~7GB)

핵심 설계: Two-pathway attention이 Pi0.5의 핵심 — image/text expert와 action expert가 동일 LLM layer 안에서 병렬로 처리됨. action_in/out_proj이 32-dim action을 LLM 공간(1024-dim)으로 project하고 flow matching denoiser로 action chunk 생성.

Action / State Spec — norm_stats.json 분석

state 32-dim + action 32-dim 모두 정규화 상수 보유. active dims (std > 1e-6) 기준:

dim	의미	std (raw)
0–5	arm OSC_POSE: Δx Δy Δz Δrx Δry Δrz	0.31–0.37 (xyz), 0.08–0.10 (rpy)
6	gripper (binary-ish)	0.99
7–9	base velocity: vx vy wz	0.14–0.23
10	(padding)	0.0
11	torso	0.70
12–31	cross-embodiment padding (unused)	1.0 (의미 없음)

확인: PandaOmron HYBRID_MOBILE_BASE의 11-dim action (arm 6 + gripper 1 + base 3 + torso 1)과 정확히 일치. cascade adapter에서 active 11 dim만 robosuite에 전달하면 됨.

Loading API — openpi 라이브러리 필요

# 예상 openpi v0.x API (다음 세션에 repo 클론 후 정확한 이름 확인 필요) from openpi.policies import policy as openpi_policy from openpi.training import config as openpi_config config = openpi_config.get_config("pi05_robocasa") policy = openpi_policy.Policy.from_checkpoint( config=config, checkpoint_dir="/home/.../pi05_pretrain_human300/multitask_learning/75000", ) # Inference action_chunk = policy.infer({ "observation/image/agentview_left": np.uint8 (256, 256, 3), "observation/image/agentview_right": ..., "observation/image/eye_in_hand": ..., "observation/state": np.float32 (32,), "prompt": "navigate to the freezer", # ← MemER subtask가 여기 들어감 }) # action_chunk: (action_horizon=H, 32). active 11 dim만 robosuite에 전달.

미검증: openpi의 정확한 API와 RoboCasa config 이름은 다음 세션에 repo 클론 + 문서 확인 필요. 위 코드는 일반적인 openpi 패턴 추정.

Cascade Adapter 골격 — MemerPi05Policy(BasePolicy)

3-venv IPC 설계:

컴포넌트	venv	역할
eval_runner orchestrator	`robocasa` (3.11)	scene-mem env 로드 + rollout 루프
MemerPi05Policy wrapper	`robocasa`	BasePolicy.reset / get_action 구현
MemER inference subprocess	`memer` (3.11, torch 2.11)	LoRA 머지 ckpt + Qwen3-VL infer
Pi0.5 inference subprocess	`robomme_policy_learning/.venv` (3.12, jax 0.5.3)	orbax/openpi action chunk emit

class MemerPi05Policy(BasePolicy): def reset(self, instruction: str, video_paths: Sequence[str]) -> None: # 1. MemER state reset self.memer.reset(instruction=instruction) # 2. prior 비디오 → MemER FIFO 채우기 (stride=100, ~80s) for cam_frames in decode_videos(video_paths[:2]): # eye_in_hand 제외 for chunk in sliding_chunks(cam_frames, stride=100): self.memer.step(*chunk) def get_action(self, env) -> np.ndarray: obs = env._get_observations() # MemER → current_subtask 갱신 memer_result = self.memer.step(obs["robot0_agentview_left_image"], obs["robot0_agentview_right_image"]) subtask = memer_result.current_subtask # Pi0.5 action chunk (IPC to JAX subprocess) action_chunk = self.pi05.infer( images={...}, state=build_proprio(obs), # 32-dim padded prompt=subtask, ) return action_chunk[0] # 32-dim; active 11만 robosuite가 사용

3 결과

아키텍처 확정 항목

Model family

PaliGemma-3B

Total params

~3.5B

Action dim

32 (11 active)

Format

orbax/JAX

ViT layers

LLM layers

비용 추정 — 257 scenario 전체 평가

단계	시간	비고
Memory build (per scenario)	~80s	3 videos × stride 100 × 0.27s/call; 1회만, test 간 공유
Per-test rollout	~200s	1500 step × (Pi0.5 ~0.1s + MemER 매 N step ~0.8s + mujoco ~0.01s)
Per-scenario	~680s (~11분)	80s + avg 3 tests × 200s
257 scenario sequential	~47h	순차 실행
8 GPU SLURM batch	~6h	sbm, 병렬화
30 scenario subset	~5h	combo_002(27개) 우선 — 권장 시작점

5개 필수 검증 체크리스트 (다음 세션 1순위)

① Subtask vocab matching

MemER LoRA 출력이 Pi0.5의 학습 분포 prompt와 정확히 매치하는지. dusting MemER는 "place X on bottom shelf"만 출력 → RoboCasa 학습된 Pi0.5에겐 의미 없음. RoboCasa LoRA (Job 2448) 학습 후 비교 필수.

② Camera obs key naming

scene-mem 비디오: robot0_agentview_left.mp4. robosuite live obs: robot0_agentview_left_image(일반적). eval_runner.py의 env._get_observations() 출력 1회 인쇄해서 정확한 키 확인 필요.

③ Camera resolution

비디오 256×256, MemER resize → 320×180. Pi0.5 학습 시 해상도는 norm_stats에 image stats 없음 → openpi config 직접 확인 필요.

④ Action chunking

Pi0 family는 보통 action_horizon=50 chunked output. eval_runner.py는 1 step = 1 action → chunk[0]만 쓰거나 매 50 step마다 chunk를 새로 받는 두 가지 전략. latency vs reactivity tradeoff.

⑤ State padding layout

state 32-dim 중 active ~12. 어느 슬롯이 어떤 proprio 값인지 정확한 매핑은 openpi/robocasa training config에서 확인. norm_stats의 dim 순서와 robosuite obs 키 순서 대조 필요.

4 Takeaway

Pi0.5는 "표준 transformer + action head" 설계 — two-pathway가 핵심

PaliGemma LLM의 동일 layer 안에 text/image expert(2048-dim)와 action expert(1024-dim)가 병렬로 co-exist하는 구조. 이 설계 덕분에 언어 prompt("navigate to the freezer")와 시각 관찰이 action denoising에 직접 조건으로 작용한다. MemER subtask string이 곧 Pi0.5의 prompt 입력 → cascade 연결점이 명확하다.

호환성 확인: robomme_policy_learning/.venv에 이미 jax 0.5.3 + flax + openpi(local copy) + orbax 가 있어 Pi0.5 추론 환경은 추가 설치 없이 재사용 가능. memer venv (torch 2.11)와 분리해서 process IPC로 연결하면 됨.

Vocab mismatch는 cascade의 최대 리스크: dusting MemER 그대로 cascade하면 Pi0.5가 이해 못 하는 subtask를 받아 score 0% 가능. RoboCasa LoRA (Job 2448) 완료 + vocab sanity check가 cascade 선행 조건.

openpi API 미검증: 위 loading 코드는 일반적 패턴 추정. 다음 세션에 openpi repo 클론 + pi05_robocasa config 이름 확인이 필수 첫 단계.

5 Next Steps

P0 — Job 2448 (RoboCasa LoRA 학습) 상태 확인

COMPLETED → swift export --merge_lora → combo_002 probe → vocab 비교 (dusting vs LoRA).
FAILED → slurms/robocasa_memer_tier1_lora_2448.out 확인. max_length=3200 부족 / OOM 예상.

P1 — openpi 라이브러리 확인 + Pi0.5 로딩 검증

robomme_policy_learning/.venv의 src/openpi/ local copy에서 Policy.from_checkpoint() API 확인. pi05_robocasa config 이름 확정. login node 1 GPU에서 1 forward pass 검증(≤100 step debug threshold).

P2 — MemerPi05Policy 어댑터 작성

memer/memer_eval/memer_pi05_policy.py — 위 골격 기반. 단 vocab sanity check 완료 후 RoboCasa LoRA로 MemER 교체. 1단계 MVP: dusting LoRA placeholder cascade, combo_002 × 1 test, env.step() 1500 step 완주만 확인.

P3 — 30 scenario subset 평가 → 257 scenario SLURM batch

combo_002 (27 scenarios) 먼저 (~5h, login node or 1-GPU sbm). memory_success baseline 측정. 그 다음 257 scenario 전체 → SLURM 8 GPU 병렬 (~6h).