PI MEM (Multi-Scale Embodied Memory) → RoboMME zero-shot API 적용 분석

본 문서는 PI Mem 논문(https://www.pi.website/download/Mem.pdf, Torne et al., 2025)을 정독한 뒤, 현재 RoboMME 평가에 쓰고 있는 MemER LoRA 파이프라인 (exp-closed_loop_full_sweep.md, avg sr 44%) 을 어떻게 PI Mem 의 long-term text summary 메커니즘으로 대체/보강할 수 있는지 정리한 research 노트.

핵심 요청: VLM 자체는 API (zero-shot) 이고 pi0.5 가 short-term을 비디오로 받지 않으므로, text summary 만으로 long-horizon memory 효과를 낼 수 있는 system prompt 설계 가 본 문서의 산출물.

배경/목적 (왜)

1. MemER LoRA 의 한계. 16 task × 50 ep sweep 결과 avg sr 44%, 일부 task 들 (PatternLock 16%, RouteStick 10%, StopCube 4%, InsertPeg 4%) 은 high-level subtask 가 부정확 한 게 주된 fail 모드로 추정됨. 또한 keyframe FIFO (1D clustering 으로 dist≤8 인 frame 들을 median 으로 압축) 는 - "어떤 subtask 가 완료 되었는가" 라는 semantic 진행 상태 를 직접적으로 표현하지 않음. 그림(이미지)으로만 들어감. - VLM 이 그 frame 들을 의도대로 해석한다는 보장이 없음 — 그냥 keyframe 의 visual 만 보고 해석.
2. PI Mem 의 통찰. (a) long-horizon = compressed natural language summary mt, (b) short-horizon = video encoder. 그리고 mt 는 high-level policy πHL 가 스스로 다음 token 으로 출력 한다. 즉 summary 갱신 자체가 모델의 행동.
3. 우리 setup 의 실용 제약. - pi0.5 (low-level πLL) 는 현재 obs + subtask string 만 받는 구조 — text summary 는 πLL 입력에 들어가지 않는다. PI 논문도 동일 (πLL 은 mt 를 받지 않고 subtask 만 받음). - 그러므로 우리가 "text summary 잘 쓰자" 의 의미는 πHL 호출 자체의 품질 (= 다음 subtask 추론 정확도) 개선 이지, πLL 를 바꾸는 게 아니다. - πHL 을 finetune 하지 않고 API VLM (Claude/Gemini/GPT-4o) zero-shot 으로 대체 → MemER LoRA 의 in-domain 가산점은 잃지만, 더 강한 reasoning + JSON 안정성 + grounded coord 출력 능력을 얻는다.
4. 이 문서의 목표: - PI Mem method 를 우리 codebase 의 어떤 module 에 어떻게 끼울지 결정. - zero-shot API VLM 으로 text summary mt + 다음 subtask + grounded <y,x> 를 안정적으로 emit 시키는 system prompt skeleton 을 설계.

작업 내용 (어떻게)

1. PI Mem 논문 핵심 정리

1.1 Probabilistic factorization (Sec III-A)

논문 식:

π(at:t+H, lt+1, mt+1 | ot−T:t, mt, g)
≈ πLL(at:t+H | ot−K:t, lt+1, g)
  πHL(lt+1, mt+1 | ot, mt, g)

• g : task goal (자연어).
• mt : 현재 시점까지의 자연어 summary ("language memory").
• lt+1 : 다음 subtask 자연어.
• πHL : low-frequency 호출, lt+1 과 mt+1 을 동시에 생성. mt 는 자기 자신의 직전 출력을 다시 받는 autoregressive 구조.
• πLL : action chunk 생성. K ≪ T (e.g., 6 frames) 의 short obs window + 새 subtask 만 받음. mt 를 받지 않음.

우리 한국어 풀이: 고수준 호출은 "지금까지 무슨 일이 있었는지를 한 줄로 다시 적고, 다음에 뭘 해야 하는지를 정한다" 의 한 step. 저수준은 그 한 줄 subtask 를 받아 행동을 만든다.

1.2 mt 갱신의 작동 예시 (Sec III-B 직접 인용)

mt:   I placed a plate in the cabinet and moved to the counter.
        ↓
mt+1: I placed a plate in the cabinet, moved to the counter, and picked up a bowl.

1.3 Compression 이 필수인 이유 (Sec III-B + Fig 6)

학습 데이터 생성: per-episode subtask 시퀀스 (with success/fail) 를 off-the-shelf LLM 에 주고 "요약하되 이후 의사결정에 더는 필요 없는 정보는 버려라" 로 prompt → 그 출력을 GT mt 시퀀스로 사용.

"instead of remembering the precise attributes of all objects … 'I put a light green bowl, a dark blue bowl and a bright yellow bowl into the top right cabinet', it is often sufficient to just remember 'I placed three bowls in the top right cabinet'"

핵심 ablation (Fig 6 'Naive Text + Video'): subtask history concat 만 하면 성능이 크게 떨어짐. 원인은 train-inference distribution shift — training data 는 사람 demo 라 subtask 가 한 번씩 밖에 나오지 않지만, inference 때는 같은 subtask 가 fail 로 반복 출력되어 "pick up bowl" × 3 같이 누적된다. compressed memory 는 성공 까지 mt 갱신을 미루므로 이 shift 가 사라진다.

→ 우리 zero-shot 설계의 1번 규칙: "실패 / 진행중인 시도는 mt 에 절대 쓰지 마라".

1.4 Short-term video memory (Sec III-C, 우리는 skip)

• 6 obs (5 past + current), stride 1s pre-train; 18 frames / 54s post-train.
• ViT 4 번째 layer 마다 spatial → causal-temporal attention 추가.
• past timestep token drop 으로 backbone 입력 token 수 = single-frame VLA 와 동일.
• 이 부분은 우리가 안 함 — pi0.5 short-context VLA 는 현재 단일 obs 만 받음. text summary 에서 short-term 정보까지 (한정적으로) 인코딩해야 함을 의미.

1.5 실험 결과 요약 (Sec IV)

→ Long-horizon task 에서는 compressed text 만으로도 video 단독과 동등 이상 의 효과. → 따라서 video 없이도 text 가 "잘" 만들어지면 long-horizon task 에서 의미있는 gain 가능 — 우리 가설.

2. 현재 RoboMME 코드 흐름 (통합 지점 식별)

examples/robomme/subgoal_predictor.py 에 4 개의 subgoal predictor 가 있음: - NullSubgoalPredictor — 사용 X. - GeminiSubgoalPredictor — 이미 zero-shot API VLM 경로. Gemini 2.5 flash-lite 기본. per-task prompt_dict_grounded[<TASK>] 로 system prompt 구성. video clip 단위로 호출. - QwenVLSubgoalPredictor — finetuned Qwen3-VL adapter (memer 가 아닌 baseline). - MemERSubgoalPredictor — Qwen3VLModelMemER 인스턴스. keyframe FIFO 메모리 + LoRA.

2.1 MemER 의 prompt + memory 흐름 (api_memer.py)

System prompt (전체):

You are a robot program that predicts actions. The current input images from
the front-view camera shows the most recent actions the robot has executed.
The past keyframes are selected frames of particular importance from all the
actions the robot has executed so far. Based on these, output the current
subtask the robot should execute and nothing else. Some tasks may have a video
input for initial setup, some may not.

Return a JSON with:
- current_subtask: the action that should be executed at the current timestep
- keyframe_positions: list of frame positions (1-indexed) from the current input
  images where actions change

User prompt 매 호출:

[The task has a video input for initial setup: <video>]   # task 별
The task goal is: {task_goal}
Here are the selected frames from the entirety of the full execution that are
of particular importance: [<image>, <image>, ...]    ← key_frame_paths (FIFO)
Here is current input image list from the front-view camera:
[<image>, <image>, ..., <image>]                     ← 최근 8 프레임 (stride 2)

What subtask should the robot execute and what is the keyframe position?

메모리 동작: - add_execution_frame(image) — 매 step PNG 저장. - _get_current_execution_frame_paths() — 마지막 idx 부터 stride 2 로 8개 frame. - update_history_subgoals(response) — VLM 이 직접 1-indexed keyframe_positions 출력 → 해당 step PNG 들이 key_frame_paths 에 들어감. - merge_key_frame_paths(dist=8) — sim-step 거리 ≤ 8 이면 같은 group, 그 group 의 median 만 남김.

→ memory 는 image 의 FIFO. 자연어 summary 가 아예 없음. VLM 이 keyframe 이미지를 보고 다시 "이게 무슨 의미였지" 를 매번 재해석해야 함.

2.2 Gemini 경로의 prompt 구조 (gemini/prompts/base.py)

이쪽은 per-task system prompt + few-shot example 로 zero-shot 동작:

SYSTEM_PROMPT = """You are a helpful assistant ...
Possible subgoals:{subgoals}
Example:{example}
Output Format:
- 첫 frame: {"subgoal_sequence": "...", "subgoal": "..."}
- video clip 들어올 때: {"description": "...", "subgoal": "..."}
"""

• {subgoals} : task별로 "- pick up the [first/second/third] cube at " 등 4-7 줄.
• {example} : 모범 시퀀스 1 ~ 5 + notes (e.g. "bin sucks cubes, count yourself").
• grounded 변형은 GROUNDED_SUBGOAL_INFORMATION 추가 ("normalized to 0-1000").

→ 이 구조에 memory_summary 필드를 추가 하기만 하면 PI Mem 의 mt 가 그대로 들어간다. per-task notes 는 그대로 살림 (e.g., "bin sucks cubes 라서 push 하지 마라").

3. 통합 설계: zero-shot API + text summary

3.1 데이터 흐름 (proposed)

[Episode start]
  - mt = ""                                      # empty memory
  - πHL call 1 (with task_goal, demo_video?, image, mt="") 
     → emits {memory_summary, subgoal, point}
  - mt ← memory_summary                          # high-level autoregressive
  - subtask + point → pi0.5 (πLL) → action chunk
  - sim 수십 step 진행 …
[Periodic high-level call] (e.g., 매 K=48 step 마다 — Gemini path 와 동일)
  - πHL call (task_goal, current image, mt) → new {memory_summary, subgoal, point}
  - mt ← new.memory_summary
  - if subgoal != prev_subgoal: switch
  - else: keep last subtask going

핵심 차이 vs MemER: - 메모리 = 자연어 한 문단 (≈30~60 token), keyframe image 아님. - VLM 호출 빈도 ↓ (매 step → 매 K step) — pi0.5 가 한 subtask 를 K step 동안 끌고감. - VLM 자체는 cloud API → in-domain finetune 없음.

3.2 system prompt skeleton (zero-shot, task-agnostic 부분)

per-task subgoals / example / notes 는 기존 prompts/<TASK>.py 그대로 재사용. 그 위에 task-agnostic memory 섹션을 얹는다:

You are the HIGH-LEVEL controller for a robotic arm. You decide ONE subtask at a
time and maintain a compressed natural-language MEMORY of progress made so far.
A separate low-level controller will execute the subtask you choose.

# Possible subgoals (task-specific)
{subgoals}

# Example episode (task-specific)
{example}

# Per-task notes
{notes}

# Memory (mt) — what to write
The memory is a 1-2 sentence summary of EVENTS THAT HAVE COMPLETED so far in
the current episode and that still affect the next decision.

Hard rules for memory:
1. Only describe COMPLETED subtasks. Never describe attempts, failures, or
   "is approaching X" states.
2. If the previous subtask in the input has not yet visibly completed in the
   current image, COPY the previous memory string UNCHANGED and continue
   issuing the same subtask.
3. Compress aggressively. Replace enumerations with counts:
   "I picked up a red cube and a blue cube" → on completion of the third pick,
   write "I have picked 3 cubes". Drop colors once they no longer affect routing.
4. Never invent events you cannot verify from the current image.
5. After the FINAL subtask of the example sequence is completed, set memory to
   "task complete" and reissue the final subtask.

# Grounded coordinates
For subgoals that include a target object, always emit a 2-D point in the
PIXEL coordinate frame normalized to 0-1000, e.g. "pick up the red cube at
<450, 612>". Use the LAST input image to ground the point.

# Output format (STRICT JSON, no prose, no markdown fences)
{
  "scratchpad": "<= 30 words of reasoning, NOT used downstream>",
  "memory_summary": "<= 60 tokens, follows the rules above>",
  "subgoal":        "<one of the subgoals listed above, with grounded <y,x>>"
}

User message 매 호출:

Task goal: {task_goal}
Previous memory: {mt or "(empty)"}
Previous subgoal: {last_subgoal or "(none)"}

Current image: <image>
[Demo video for initial setup: <video>]   # task 별, 첫 호출 한 번만

3.3 핵심 prompt 설계 결정 + 근거

3.4 어떤 task 가 text-only 로 충분한가 (가설)

PI 논문 ablation (Fig 6) 의 Only Text Memory 가 video-only 와 거의 동등 → text-only 로 충분한 task 의 조건은 (a) 진행 상태가 이산적 이고 (b) 실패 횟수 보다 완료 횟수 가 결정적인 경우.

현재 RoboMME 16 task 분류:

→ counting/sequencing tier 에 가장 큰 즉시 gain. precision tier 는 grounding 개선이 별도 작업.

3.5 비교 표 — MemER LoRA vs PI-Mem-style zero-shot API

4. 우려 / 근거

1. Latency. API call ~1-3 s/call → 50 ep × ~20 call/ep = 1000 call/task ≈ 0.5~1 시간/task. 16 task 면 8~16 hr. SLURM 영향 없음 (GPU 안 씀). Anthropic / Gemini rate limit 만 주의.
2. JSON 안정성. 큰 API 모델은 strict JSON 잘 출력하지만 0.5-1% 실패 가능. Gemini path 의 parse_markdown_json (코드 fence 까지 처리) 재사용 + retry-once on parse fail.
3. mt drift. 매 호출 fresh context 로 mt 를 재투입하므로 "memory of memory" drift 가능. PI 논문은 학습으로 mt 의 정확한 형태를 강제하지만, 우리는 system prompt rule 만으로 강제 — 검증 필요. 첫 실험에서 mt log 다 저장해 사람이 한 번 훑어봐야 함.
4. Grounded coord 정확도. Qwen3-VL LoRA (in-domain) vs API (OOD) 의 격차가 클 수 있음. PI 의 πHL 은 grounding 안 함 (subtask string only). 우리는 grounding 도 API 가 해야 함 → MoveCube/InsertPeg 같은 precision-critical task 에서 risk.
5. Demo video 활용. 첫 호출에서만 video 를 input 으로 주고, 이후엔 mt 의 "I observed the demo: X→Y→Z" 줄만 본다. video 재투입 비용 절감 + summary 에 demo 의도가 명시적으로 박힘.

결과 (수치 — 가설 phase, 실측 다음 단계)

이 단계에서 측정 한 수치는 없음. 대신:

1. Mem 논문 ablation (Fig 6) 에서 Only Text Memory ≈ Only Video Memory, 둘 다 No Memory 대비 큰 상승 (Average 기준 ~30→55 progress). Naive concat 은 도리어 떨어짐. ⇒ "compressed text 만으로도 의미있는 효과" 는 검증된 사실.
2. 현 MemER LoRA: 16 task avg sr 44%. counting tier (BinFill 64%, PickXtimes 76%) 와 sequencing tier (PatternLock 16%, ButtonUnmaskSwap 18%) 의 격차 큼 — 후자에서 가장 큰 gain 여지.
3. Gemini path 기존 동작: leaderboard 비교 미수치 (별도 실험 필요). 본 제안의 baseline 으로 사용 예정.

의미 (Takeaway)

1. PI Mem 의 핵심은 video encoder 가 아니라 compressed text mt 의 autoregressive 출력 이다. video 부분을 떼고 text 만 가져와도 long-horizon task 에서 의미있는 효과.
2. 우리 codebase 는 이미 절반 됨. Gemini path 가 zero-shot API VLM 인프라 (per-task prompt + JSON parse + grounded coords) 를 다 가지고 있음. memory_summary 한 필드 추가 + system prompt 규칙만 더하면 된다.
3. MemER FIFO 는 ablation 의 "Only Video Memory" 와 가장 가깝다. PI 논문 ablation 대로라면 거기에 compressed text 를 더한 게 더 좋다. 우리는 video 가 없으니 compressed text 단독 으로 가야 한다.
4. Distribution-shift rule (rule 2 = 진행중이면 mt 복사) 가 zero-shot 에서 가장 risky 한 부분. 학습으로 강제하던 걸 system prompt 로 바꾸는 셈이라, 첫 실험에서 mt 를 다 logging 해서 사람이 검수해야 한다.
5. counting/sequencing tier (BinFill/PickXtimes/PatternLock/ButtonUnmaskSwap) 가 1차 검증 target 으로 가장 적합. precision tier (InsertPeg/MoveCube) 는 grounding 정확도 문제라 따로 봐야 함.

보완점 / 다음 (Next)

Phase 0 (이번 session 외)

• API 키 확인. ~/env/api_keys.txt (read-only) 로 ANTHROPIC_API_KEY / GEMINI_API_KEY 존재 여부만 확인 — 직접 출력/전송 금지.

Phase 1 — 인프라

• subgoal_prediction/api_mem/ 신설. 기존 gemini/ 를 base 로 fork:
• prompts/base.py 에 MEMORY_RULES 블록 + OUTPUT_FORMAT_WITH_MEMORY 추가.
• api.py 에 ClaudeMemModel / GeminiMemModel 두 개 (provider switch). 첫 실험은 모델 1개로.
• subgoal_predictor.py 에 ApiMemSubgoalPredictor 클래스 추가. mt 를 episode state 에 보관, 매 call 마다 user prompt 에 주입.
• 새 module 은 git tracked, 기존 gemini/ qwenvl/ 는 무수정 (closed-loop sweep 재현성 보존).

Phase 2 — 1-task smoke test

• BinFill ep0 1 episode 로 mt log 를 사람이 검수.
• mt 가 rule 2 (진행중이면 unchanged) 를 지키는가?
• count compression 이 발동하는가? ("I have placed 1 red cube" → "2 cubes" → "3 cubes")
• JSON parse 실패율?
• grounded coord 의 OOD 격차 (vs MemER LoRA in-domain) 측정.

Phase 3 — counting/sequencing tier sweep

• BinFill, PickXtimes, SwingXtimes, ButtonUnmask, ButtonUnmaskSwap, PatternLock 6 task × 10 ep 로 빠른 ablation:
• (a) MemER LoRA (현 baseline)
• (b) API zero-shot, no memory (system prompt 에서 mt 부분 제거)
• (c) API zero-shot, with mt (본 제안)
• (d) API zero-shot, naive concat memory (rule 2 제거 — distribution-shift 가설 검증용)
• 가설: (c) > (a) on PatternLock/ButtonUnmaskSwap, (c) ≥ (a) on BinFill/PickXtimes. (d) < (c) 면 PI 논문의 distribution-shift 주장 zero-shot 에서도 재현됨.

Phase 4 — 16 task full sweep + visualization

• counting tier 에서 검증되면 16 task 전체.
• render_with_memory.py 를 mt 표시용으로 확장 — keyframe FIFO strip 자리에 mt 의 marquee 표시.

미해결 question

1. demo video task 에서 첫 호출의 <video> 입력을 mt 1 줄로 압축할 때, 그 1 줄이 정확한가 를 검증할 method 가 필요. (사람 검수 외에는 없음 → Phase 2 에서 BinFill 외에 VideoUnmask ep0 도 같이 검수.)
2. πLL (pi0.5) 는 본 mme_vla_suite checkpoint 라 symbolic-grounded-subgoal 형식의 subtask string 에 fit 되어 있음. API 가 generated subtask string 의 표현이 학습 분포와 어긋나면 πLL 이 못 따라옴. → API 출력을 반드시 per-task subgoals 리스트의 정확한 wording 으로 강제 (system prompt rule + JSON enum-like constraint).
3. mt 의 최대 길이 — 너무 길면 cost / drift, 너무 짧으면 정보 손실. Phase 2 에서 60-token soft cap, 100-token hard cap 으로 시작하고 BinFill ep0 mt 평균 길이로 fine-tune.

참고 파일 / 인용

• 논문: https://www.pi.website/download/Mem.pdf (./tmp/mem.pdf 로컬 캐시).
• 블로그: https://www.pi.website/research/memory.
• 본 repo:
• examples/robomme/subgoal_prediction/qwenvl/api_memer.py — 현 MemER inference.
• examples/robomme/subgoal_prediction/gemini/api.py — 기존 zero-shot API 인프라 (재사용 base).
• examples/robomme/subgoal_prediction/gemini/prompts/base.py — system prompt template.
• examples/robomme/subgoal_prediction/gemini/prompts/BinFill.py — per-task prompt 예시.
• examples/robomme/subgoal_predictor.py — 4 종 predictor wrapper.
• 직전 작업: 260430/exp-closed_loop_full_sweep.md — MemER LoRA sweep 결과.
• 환경: 260428/setup-robomme_b200_cluster.md.