← Back
2026-04-30 · plan

PI Mem-style zero-shot API πHL 구현 계획 (`api_mem` module)

TL;DR. > 본 문서는 [research-pi_mem_text_summary_adaptation.md](research-pi_mem_text_summary_adaptation.md)

PI Mem-style zero-shot API πHL 구현 계획 (api_mem module)

본 문서는 research-pi_mem_text_summary_adaptation.md 의 Phase 1~4 를 코드 단위까지 분해한 implementation plan. 한 명이 처음부터 따라가도 막힘 없이 BinFill ep0 까지 도달 가능하도록 파일/함수/JSON schema/CLI flag 까지 명시.

원칙 (이전 작업의 educated lesson): - 기존 gemini/, qwenvl/ 경로는 무수정. fork 가 아니라 새 모듈로. - 모든 신규 코드는 git tracked, tmp/ 는 임시 산출물 한정. - try/except 로 에러 숨기지 않음. 첫 1 ep 검증 단계까지는 raise + log. - login GPU 안 씀 (API call only). SLURM 도 일단 안 씀 (90 min 안에 1 task 끝나는 빈도면).

Phase 0 — 환경 / 키 (DONE 2026-04-30)

실제 파일 위치 (글로벌 CLAUDE.md 와 다름)

결과

Phase 1 의 loader 수정안

def _load_gemini_key(self):
    p = Path.home() / "envs" / "api_keys.txt"     # NOTE: "envs" (plural)
    for line in p.read_text().splitlines():
        if ":" not in line:
            continue
        name, value = line.split(":", 1)
        if name.strip().lower() == "gemini":
            os.environ["GEMINI_API_KEY"] = value.strip()
            genai.configure(api_key=os.environ["GEMINI_API_KEY"])
            return
    raise RuntimeError("Gemini key not found in ~/envs/api_keys.txt")

Phase 1 — 인프라 (api_mem 모듈 신설, 1 day)

1.1 파일 트리

examples/robomme/subgoal_prediction/api_mem/
├── __init__.py              # 빈 파일
├── api.py                   # ApiMemModel (Gemini provider 1개부터)
├── prompts/
│   ├── __init__.py          # prompt_dict_grounded_with_memory 빌드
│   ├── base.py              # MEMORY_RULES + OUTPUT_FORMAT_WITH_MEMORY + helper
│   └── <TASK>.py × 16       # 기존 gemini/prompts/<TASK>.py 의 subgoals/example/notes 를
│                            # base.py 의 새 template 로 다시 fmt
└── README.md                # 모듈 사용법

1.2 prompts/base.py — system prompt 설계

기존 gemini/prompts/base.py 의 SYSTEM_PROMPT 와 분리해서 새 template 정의.

# api_mem/prompts/base.py

GROUNDED_COORD_INFO = """\
For subgoals that include a target object, emit a 2-D point in the PIXEL
coordinate frame normalized to 0..1000, e.g. "pick up the red cube at <y, x>".
Use the LAST input image to ground the point. If the object is not visible in
the last image, copy the coordinate from the previous subgoal."""

MEMORY_RULES = """\
You also maintain a compressed natural-language MEMORY of the events that
have happened in this episode so far.

Hard rules for the 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 re-issue the
   same previous subgoal (not a new one).
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 attributes once they no longer affect
   the next decision.
4. Never invent events you cannot verify from the current image.
5. After the FINAL subtask in the example sequence is completed, set memory to
   "task complete" and re-issue the final subgoal."""

OUTPUT_FORMAT_WITH_MEMORY = """\
Output STRICT JSON only, no prose, no markdown fences:
```json
{
  "scratchpad":      "<= 30 words of reasoning, ignored downstream>",
  "memory_summary":  "<= 60 tokens, follows the memory rules above>",
  "subgoal":         "<one of the listed subgoals, with grounded <y, x>>"
}
```"""

SYSTEM_PROMPT_WITH_MEMORY = """\
You are the HIGH-LEVEL controller of a robotic arm. At each call you decide
ONE subgoal for a separate low-level controller to execute, and you maintain
a compressed MEMORY of progress so far.

# Possible subgoals (this task)
{subgoals}

# Example episode (this task)
{example}

# Per-task notes
{notes}

# Memory rules
""" + MEMORY_RULES + """

# Coordinate grounding
""" + GROUNDED_COORD_INFO + """

# Output format
""" + OUTPUT_FORMAT_WITH_MEMORY


USER_PROMPT_FIRST_CALL = """\
Task goal: {task_goal}
Previous memory: (empty)
Previous subgoal: (none)

Current image follows.
{demo_video_marker}"""

USER_PROMPT_STEP = """\
Task goal: {task_goal}
Previous memory: {mt}
Previous subgoal: {last_subgoal}

Current image follows."""

DEMO_VIDEO_MARKER = "Demo video for initial setup follows after the image. " \
                    "Use it to learn the desired sequence; do NOT re-watch on later calls."

{notes} 는 기존 gemini/prompts/<TASK>.py 의 notes 변수 그대로 재사용. {subgoals} / {example} 는 grounded 변형 (subgoals_grounded, example_grounded) 사용.

1.3 prompts/<TASK>.py × 16

기존 gemini/prompts/<TASK>.py 의 변수 4 개 (subgoals_grounded, example_grounded, notes, optionally keep_period_hints) 만 import 해서 새 SYSTEM_PROMPT 에 fmt:

# api_mem/prompts/BinFill.py
from subgoal_prediction.gemini.prompts.BinFill import (
    subgoals_grounded, example_grounded, notes
)
from subgoal_prediction.api_mem.prompts.base import SYSTEM_PROMPT_WITH_MEMORY

BinFill_SYSTEM_PROMPT_API_MEM = SYSTEM_PROMPT_WITH_MEMORY.format(
    subgoals=subgoals_grounded,
    example=example_grounded,
    notes=notes,
)

→ per-task 변경 시 gemini/prompts/<TASK>.py 만 고치면 양쪽 다 반영. fork 하지 않음.

1.4 api.py — ApiMemModel

# api_mem/api.py

import os, re, json, time
from pathlib import Path
from typing import Optional
import imageio.v3 as iio
import google.generativeai as genai

from subgoal_prediction.api_mem.prompts import prompt_dict_with_memory


class ApiMemModel:
    """Zero-shot API VLM that emits {memory_summary, subgoal, scratchpad}.

    Per-call lifecycle:
      start_new_episode(save_dir, demo_video, task_goal, task_id)
      get_subgoal(current_image, last_subgoal_was_complete) -> subgoal_str
      end_episode()
    """

    def __init__(self, provider: str = "gemini", model_name: Optional[str] = None,
                 image_size: tuple = (256, 256)):
        self.provider = provider
        self.image_size = image_size

        if provider == "gemini":
            self._load_gemini_key()
            self.model_name = model_name or "gemini-2.5-pro"
        else:
            raise NotImplementedError(f"Provider {provider} not yet supported")

    def _load_gemini_key(self):
        key_file = Path.home() / "env" / "api_keys.txt"
        for line in key_file.read_text().splitlines():
            if line.startswith("GEMINI_API_KEY="):
                os.environ["GEMINI_API_KEY"] = line.split("=", 1)[1].strip()
                genai.configure(api_key=os.environ["GEMINI_API_KEY"])
                return
        raise RuntimeError("GEMINI_API_KEY missing in ~/env/api_keys.txt")

    def start_new_episode(self, save_dir: str, demo_video, task_goal: str, task_id: str):
        self.save_dir = save_dir
        os.makedirs(save_dir, exist_ok=True)
        ep_name = os.path.basename(save_dir)
        self.log_path = os.path.join(os.path.dirname(save_dir), f"{ep_name}_ApiMem_log.jsonl")

        self.task_goal = task_goal
        self.task_id = task_id
        self.mt = ""                     # language memory
        self.last_subgoal = None
        self.last_subgoal_completed = True
        self.uploaded_files = []         # for cleanup
        self.call_idx = 0

        # Build system instruction (per-task)
        sys_prompt = prompt_dict_with_memory[task_id]
        self.model = genai.GenerativeModel(
            model_name=self.model_name,
            system_instruction=sys_prompt,
        )

        # Cache demo video upload (one-shot, reused on first call)
        if demo_video is not None and len(demo_video) > 0:
            video_path = os.path.join(save_dir, "demo_video.mp4")
            iio.imwrite(video_path, demo_video, plugin="pyav", codec="h264", fps=30)
            uf = genai.upload_file(path=video_path)
            self.uploaded_files.append(uf)
            while uf.state.name == "PROCESSING":
                time.sleep(0.2); uf = genai.get_file(uf.name)
            self.demo_video_file = uf
        else:
            self.demo_video_file = None

    def get_subgoal(self, current_image, last_subgoal_completed: bool) -> str:
        # 1. Save current image to disk + upload
        img_path = os.path.join(self.save_dir, f"step_{self.call_idx}_image.png")
        iio.imwrite(img_path, current_image)
        img_uf = genai.upload_file(path=img_path)
        self.uploaded_files.append(img_uf)

        # 2. Build user prompt
        if self.call_idx == 0:
            user_text = USER_PROMPT_FIRST_CALL.format(
                task_goal=self.task_goal,
                demo_video_marker=DEMO_VIDEO_MARKER if self.demo_video_file else ""
            )
            parts = [user_text, img_uf]
            if self.demo_video_file:
                parts.append(self.demo_video_file)
        else:
            user_text = USER_PROMPT_STEP.format(
                task_goal=self.task_goal,
                mt=self.mt or "(empty)",
                last_subgoal=self.last_subgoal or "(none)"
            )
            parts = [user_text, img_uf]

        # 3. Call API. retry once on parse fail.
        for attempt in range(2):
            response = self.model.generate_content(parts)
            response_text = response.text
            parsed = self._parse_json(response_text)
            if parsed is not None:
                break
            print(f"[api_mem] parse fail (attempt {attempt+1}), response was:\n{response_text}")
        else:
            raise RuntimeError("API output failed to parse twice in a row.")

        # 4. Update mt (rule 2: only update if last subgoal completed)
        new_mt = parsed["memory_summary"]
        new_subgoal = parsed["subgoal"]

        # Log raw
        self._log_call(user_text, response_text, parsed)

        # Trust the model's rule-2 compliance (system prompt enforces).
        self.mt = new_mt
        self.last_subgoal = new_subgoal
        self.call_idx += 1

        return new_subgoal

    def _parse_json(self, text: str) -> Optional[dict]:
        # strip markdown fences
        text = text.strip()
        if text.startswith("```"):
            text = re.sub(r"^```(?:json)?\s*", "", text)
            text = re.sub(r"\s*```$", "", text)
        try:
            obj = json.loads(text)
        except json.JSONDecodeError:
            return None
        for k in ("memory_summary", "subgoal"):
            if k not in obj:
                return None
        return obj

    def _log_call(self, user_text, response_text, parsed):
        with open(self.log_path, "a") as f:
            json.dump({
                "call_idx": self.call_idx,
                "user_text": user_text,
                "response_text": response_text,
                "parsed": parsed,
                "prev_mt": self.mt,
                "prev_subgoal": self.last_subgoal,
            }, f)
            f.write("\n")

    def end_episode(self):
        for uf in self.uploaded_files:
            try: genai.delete_file(uf.name)
            except Exception: pass
        self.uploaded_files = []

1.5 subgoal_predictor.py — ApiMemSubgoalPredictor

# 추가 import
from subgoal_prediction.api_mem.api import ApiMemModel


class ApiMemSubgoalPredictor(SubgoalPredictorBase):
    def setup_api(self) -> None:
        self.api = ApiMemModel(
            provider=self.args.api_mem_provider,
            model_name=self.args.api_mem_model_name,
        )
        print(f"[robomme] ApiMem ({self.args.api_mem_provider}) agent setup")

    def start_episode(self, epstate, env_runner):
        super().start_episode(epstate, env_runner)
        self.episode_dir = os.path.join(self.save_dir, self.env_name, f"ep{self.episode_id}")
        # demo_video = epstate.image_buffer[:-1] if task has demo, else None
        demo = epstate.image_buffer[:-1] if self.env_name in TASK_WITH_VIDEO_DEMO else None
        self.api.start_new_episode(
            save_dir=self.episode_dir,
            demo_video=demo,
            task_goal=self.task_goal,
            task_id=self.env_name,
        )
        self._call_period = 48     # match Gemini path
        self._step_count = 0

    def step(self, epstate):
        self._step_count += 1

    def get_subgoal(self, count, current_subgoal, last_subgoal):
        if count > 0 and count % self._call_period != 0:
            return current_subgoal, False
        # Get latest image from epstate via env_runner (already in api state)
        # NOTE: we need image via predictor api — pass from caller.
        # Actually current_image comes via step() in epstate.image_buffer[-1]
        last_img = self.env_runner.last_image  # add accessor or store in step()
        subgoal = self.api.get_subgoal(last_img, last_subgoal_completed=True)
        return subgoal, False

    def end_episode(self, epstate, success_flag):
        self.api.end_episode()

1.6 subgoal_predictor.py build switch

def build_subgoal_predictor(args, save_dir):
    if args.use_gemini:    return GeminiSubgoalPredictor(args, save_dir)
    if args.use_qwenvl:    return QwenVLSubgoalPredictor(args, save_dir)
    if args.use_memer:     return MemERSubgoalPredictor(args, save_dir)
    if args.use_api_mem:   return ApiMemSubgoalPredictor(args, save_dir)   # NEW
    if args.use_oracle:    return OracleSubgoalPredictor(args, save_dir)
    return NullSubgoalPredictor(args, save_dir)

1.7 CLI flags (examples/robomme/eval.py 의 args)

기존 args struct 에 추가 (eval.py 무수정 원칙이지만 flag 추가는 OK, monkey-patch 와 다른 범주):

--args.use-api-mem            : bool, default False
--args.api-mem-provider       : str, default "gemini"
--args.api-mem-model-name     : str, default None  (provider default 사용)
--args.api-mem-call-period    : int, default 48

만약 eval.py 무수정을 강제로 지키려면 → scripts/run_eval_with_mplib_patch.py 에서 args 객체에 monkey-patch 형태로 attribute 주입. 하지만 cleaner 한 건 eval.py 에 attr 추가 (MemER 추가 때도 추가했음). 결정: 추가.

1.8 Files touched (Phase 1)

1.9 Success criterion

cd ~/repos/Robotics/robomme_policy_learning
python3 -c "from examples.robomme.subgoal_prediction.api_mem.api import ApiMemModel; \
            m = ApiMemModel(provider='gemini'); print('ok')"

Phase 2 — 1-task / 1-ep smoke test (BinFill ep0, 0.5 day)

2.1 실행

2.2 산출물 검수 항목 (사람이 직접 본다)

2.3 Failure modes 체크리스트 (사전)

증상 1차 의심 대응
JSON 깨짐 system prompt 의 strict-mode 약함 retry once 후 system prompt 끝에 Output strict JSON or you fail. 추가
mt 가 매 call 마다 새로 씀 rule 2 무시 system prompt 에 "If unchanged, copy verbatim" 강조 + previous memory: 를 user msg 가장 위에
count compression 안 됨 example 부족 per-task example_grounded 를 count 누적 형태로 다시 작성
좌표 OOD API 의 grounding 가 256×256 에서 약함 입력 image 를 1024×1024 로 upscale (norm 0..1000 system prompt 와 매칭)
latency > 5 s/call gemini-2.5-pro 의 video 첨부 cost demo video 첨부 첫 1 회만, 이후 image only — 이미 설계에 반영

2.4 Success criterion

Phase 3 — counting/sequencing tier ablation (1.5 day)

3.1 비교군

ID Predictor Memory 용도
A MemERSubgoalPredictor (LoRA) keyframe FIFO baseline (이미 sweep 결과 있음)
B ApiMemSubgoalPredictor provider=gemini, mt 비활성 없음 API 모델 자체 capability
C ApiMemSubgoalPredictor provider=gemini, mt 활성 compressed text 본 제안
D ApiMemSubgoalPredictor provider=gemini, naive concat concat last N subgoals distribution-shift 검증

B, D 는 코드 path 동일, system prompt 만 다르게. → --args.api-mem-prompt-variant flag 로 switch: - mem (default, 본 제안) - no_mem (memory rules 섹션 제거, JSON 에 memory_summary 빼고 subgoal 만) - naive_concat (memory_summary 자리에 직전 N=10 호출의 subgoal 을 concat 으로 user msg 에 박음, model 은 subgoal 만 출력)

3.2 Tasks (6 종, 각 10 ep)

BinFill, PickXtimes, SwingXtimes, ButtonUnmask, ButtonUnmaskSwap, PatternLock. 선정 근거: counting/sequencing tier — research-pi_mem_text_summary_adaptation.md Phase 3 의 가설 검증 대상.

3.3 Run plan

3.4 산출물

3.5 가설 (검증 target)

3.6 Success criterion

Phase 4 — 16-task full sweep + 시각화 (1 day)

4.1 실행

4.2 시각화

4.3 산출물

4.4 Success criterion

위험 / 사전 대응

위험 영향 사전 대응
Gemini API rate-limit (RPM) sweep 중단 1 task 당 sleep 200 ms 삽입, sbm retry 5회. 또는 Anthropic provider 추가.
비용 폭주 budget overrun Phase 2 에서 1 ep 당 token cost 측정 후 16 task × 50 ep × N call 곱해서 사전 estimate. budget cap (e.g. $50) 넘으면 stop.
pi0.5 가 새 subtask 표현 못 받아드림 sr 폭락 system prompt rule "한 글자도 변형 말고 listed subgoals 중 하나만" 강조, 첫 1 ep 검수에서 subtask string 의 distribution match 확인.
eval.py 변경에 의한 sweep 재현성 손상 기존 결과와 비교 불가 eval.py 의 변경은 flag 추가 한정 (default False), 기존 path 동작 무영향.
Demo video 의 첫 호출 cost 가 큼 latency 5+ s upload 한 번만, genai.upload_file 결과 episode 내 reuse — 이미 설계 반영.
API JSON parse fail 호출 실패 retry 1 회. 두 번 다 실패 시 last_subgoal 유지하고 mt unchanged → graceful degradation.

일정 / 마일스톤

Phase 작업 소요 Deliverable
0 키 확인 30 min boolean check
1 api_mem/ 모듈 + predictor + flag 1 day import 통과
2 BinFill ep0 smoke 0.5 day mt log 검수 + ep0 success
3 counting tier 4-way ablation 1.5 day 4×6 표, 가설 ≥ 2 개 검증
4 16 task full + 시각화 1 day full sr 표 + with_mt 비디오

총 ≈ 4 day (사람 검수 시간 별도 0.5 day).

Open question (이 plan 으로도 못 정한 것)

  1. Provider 선택. Gemini 2.5 Pro vs Claude Opus 4.7 — Phase 3 ablation 으로 한 번 더 비교해야 함. cost / latency / JSON 안정성 / grounding 정확도.
  2. call_period = 48 step 의 적정성. PatternLock 같이 빠른 sequence 가 있는 task 에서 너무 느릴 수 있음. Phase 2 BinFill 검수 후 task-별 hint dict 가 필요할 수도.
  3. last_subgoal_completed flag 의 신호. 본 설계는 model 에 자체 판단 떠넘김 (rule 2). 만약 model 이 자주 잘못 판단하면, sim 의 subgoal-success-detector (있는 경우) 에서 boolean 을 가져와 user prompt 에 명시 입력해야 할 수도. 1차는 model 자체 판단으로 시작.
  4. mt 의 hard cap — 60 token soft, 100 hard 로 설계했지만 실제 token 분포 확인 필요. 첫 1 ep 후 결정.

참고 / 연결