Module 5 · Project — VLA on your laptop

5C.1 · What you'll build

SmolVLA가 gym-aloha(MuJoCo)의 양팔 로봇을 구동 → 결과로 rollout.mp4(팔이 움직이는 영상) + 스텝별 로그(지시와 VLA가 낸 action 벡터)를 남깁니다. 루프를 눈에 보이게 만드는 게 목적이에요.

5C.2 · Scaffold — 파일 4개가 전부

robot.zip을 풀면 딱 4개예요. 무겁게 느껴져도 실제 프로젝트 코드는 작습니다 — 무거운 건 다운로드되는 모델·시뮬 에셋이지 우리 코드가 아니에요.

5C.3 · The stack & why — 왜 이 조합인가

다섯 저장소를 붙이는 대신 LeRobot 하나가 정책·데이터 규약·체크포인트를 묶어줘요. 모델은 SmolVLA(~450M), 시뮬은 gym-aloha(MuJoCo), 런타임은 PyTorch CPU.

5C.4 · The loop — 핵심 아이디어

VLA의 본질을 한 루프로: 카메라 이미지 + 지시 → VLA → action 벡터 → 시뮬이 로봇을 움직임 → 다음 프레임 → 반복. run_demo.py가 정확히 이 루프예요(실질 ~40줄).

🔵 점(dot)이 이 네 노드를 영원히 도는 게 이 프로젝트의 전부예요 — OBS → VLA → ACT → SIM → OBS. "루프가 엔진"이라는 걸 눈으로 붙잡아 두세요. 아래 5C.5는 이 루프를 실제 함수 이름으로 펼칩니다.

5C.5 · main() 호출 순서 — 흐름을 함수 이름으로 ★

위 루프를 run_demo.py의 실제 함수 호출 순서로 펼친 지도예요. 앞쪽 4개는 한 번만 하는 셋업, 뒤쪽은 매 스텝 도는 루프 — 초록 점이 그 루프를 계속 돕니다.

셋업 네 노드(LOADC→PRET→PROC→RESET)의 실제 코드는 이렇게 짧아요. 두 군데를 눈여겨보세요 — ① policy.reset()이 비우는 action-chunk 큐, ② 전·후처리기를 만들 때의 device 오버라이드.

    SmolVLAPolicy = load_smolvla_class()
    policy = SmolVLAPolicy.from_pretrained(args.checkpoint)
    policy.config.device = args.device
    policy.to(args.device)
    policy.eval()
    policy.reset()  # clears the internal action-chunk queue

    # The pre-processor tokenizes the instruction + normalizes inputs; the
    # post-processor un-normalizes the predicted action. These live alongside the
    # checkpoint on the hub, so we load them from the same path.
    #
    # The saved processor config hard-codes device 'cuda' — its DeviceProcessorStep
    # asserts CUDA at construction time, which fails on a Mac. Override it to our
    # device via the factory's per-step overrides (keyed by the step name).
    device_override = {"device_processor": {"device": args.device}}
    preprocessor, postprocessor = make_pre_post_processors(
        policy.config,
        pretrained_path=args.checkpoint,
        preprocessor_overrides=device_override,
        postprocessor_overrides=device_override,
    )

    env = gym.make(env_id, obs_type="pixels_agent_pos", render_mode="rgb_array")
    obs, _ = env.reset(seed=0)

5C.6 · 먼저 실행: check_setup.py

로봇-러닝 생태계는 빨리 변해요(import 경로·env id·키가 드리프트). 그래서 롤아웃 전에 스펙을 출력해 어긋남을 먼저 드러냅니다.

# check_setup.py — 4단계로 "지금 내 환경의 진짜 스펙"을 찍어줌
# 1) torch / lerobot import 되나
# 2) gym-aloha 가 env 를 registry 에 등록했나  → 등록된 id 목록 출력
# 3) env 만들어 스펙 출력: action_space.shape (정책이 내야 할 action dim!)
#                          카메라 키, agent_pos dim
# 4) SmolVLA 로드 → 체크포인트가 '기대하는' 스펙 출력:
#      cfg.input_features   (카메라 키 · state dim)
#      cfg.output_features  (action dim — 학습 당시 로봇 기준)

진짜 main()의 네 단계는 이렇게 생겼어요 — 각 단계가 print로 자기 결과를 찍어 어디서 어긋났는지 눈에 보이게 합니다.

def main() -> int:
    print("\n1. Core imports")
    # … import torch, lerobot

    print("\n2. Simulator (gym-aloha / MuJoCo)")
    import gym_aloha  # noqa: F401  (registers the gym_aloha/* envs)
    env_ids = [s for s in gym.registry.keys() if "aloha" in s.lower()]

    print("\n3. Make an env and show its spec")
    env = gym.make(env_id, obs_type="pixels_agent_pos", render_mode="rgb_array")
    ok(f"action_space: {env.action_space.shape}  (this is the # of action dims the policy must output)")
    ok(f"agent_pos dim: {obs['agent_pos'].shape}")

    print("\n4. Load the SmolVLA policy and show what IT expects")
    policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
    cfg = policy.config
    in_feats = getattr(cfg, "input_features", None)
    out_feats = getattr(cfg, "output_features", None)
    print(f"     input_features:  {in_feats}")
    print(f"     output_features: {out_feats}")

왜 중요하냐면 — 4단계가 찍는 체크포인트의 action dim(output_features)과 3단계의 시뮬 action dim(action_space.shape)이 다를 수 있어요. check_setup.py는 바로 이 두 숫자를 나란히 찍고 "다르면 다른 로봇 몸체용 체크포인트"라는 노트까지 출력합니다 — 그 불일치가 곧 embodiment 교훈의 씨앗입니다.

5C.7 · build_observation — obs 조립

매 스텝, 시뮬이 준 raw 데이터를 체크포인트가 기대하는 모양으로 맞춰요. 핵심: 모양을 하드코딩하지 않고 policy.config.input_features에서 읽어와서 맞춥니다. 아래 다이어그램이 함수 안 흐름의 지도, 그 아래가 실제 코드예요.

def build_observation(policy, obs: dict, instruction: str):
    # … read expected keys/shapes from config.input_features; fit the sim to them
    import torch
    feats = policy.config.input_features  # {key: PolicyFeature(type, shape)}
    visual_keys = [k for k in feats if k.startswith("observation.images")]
    pixels = obs.get("pixels")
    frames = list(pixels.values()) if isinstance(pixels, dict) else [pixels]
    batch = {}
    for i, key in enumerate(visual_keys):
        frame = frames[i] if i < len(frames) else frames[-1]
        _, h, w = feats[key].shape  # (C, H, W)
        img = torch.from_numpy(frame).float().permute(2, 0, 1) / 255.0  # (C,H,W), unbatched
        img = torch.nn.functional.interpolate(
            img.unsqueeze(0), size=(h, w), mode="bilinear", align_corners=False
        ).squeeze(0)
        batch[key] = img
    if "observation.state" in feats:
        want = feats["observation.state"].shape[0]
        s = np.asarray(obs["agent_pos"], dtype=np.float32)
        s = np.pad(s, (0, want - len(s)))[:want] if len(s) < want else s[:want]
        batch["observation.state"] = torch.from_numpy(s)  # (D,), unbatched
    batch["task"] = instruction  # the language half of "vision-language-action"
    return batch

5C.8 · The rollout loop — 예측 한 스텝

조립한 raw obs를 전처리기(토크나이즈+정규화+배치+device)에 통과 → select_action → 후처리기(역정규화) → 시뮬 step. 이게 매 틱 반복돼요 — 초록 점이 도는 사이클이 곧 아래 코드의 한 바퀴입니다.

for t in range(args.steps):
    raw = build_observation(policy, obs, instruction)
    batch = preprocessor(raw)                 # tokenize + normalize + batch + device
    with torch.inference_mode():
        action = policy.select_action(batch)  # (1, action_dim), normalized
    action = postprocessor(action)            # un-normalize, back to cpu
    action = np.asarray(action.squeeze(0).detach().cpu(), dtype=np.float32)

    # Embodiment adaptation: the base checkpoint outputs actions for a
    # different robot than this sim, so the dims won't match. We force-fit
    # (truncate / zero-pad) and clip to the sim's range so the LOOP runs. …
    if action.shape[0] != env_dim and not warned:
        print(…)                              # one-time mismatch warning (see 5C.9)
        warned = True
    if action.shape[0] > env_dim:
        action = action[:env_dim]
    elif action.shape[0] < env_dim:
        action = np.pad(action, (0, env_dim - action.shape[0]))
    action = np.clip(action, env.action_space.low, env.action_space.high).astype(np.float32)

    obs, reward, terminated, truncated, _ = env.step(action)
    frames.append(env.render())

5C.9 · Embodiment & 정직한 한계

루프는 끝까지 돌지만 동작은 의미가 없어요. base 체크포인트는 SO-101(단일 저가 팔)용인데 시뮬은 양팔 ALOHA라, 카메라·state·action의 모양과 의미가 다릅니다. 그 간극이 바로 embodiment 수업이에요.

이 force-fit이 실제로 뭘 하는지 눈으로 보면 이래요 — 모델이 낸 벡터를 시뮬이 원하는 슬롯 수에 억지로 맞춥니다:

위 숫자(모델 6 · 시뮬 14)는 대략이에요 — SO-101 단일 팔(≈6) 정책을 양팔 ALOHA(14) 시뮬에 넣으면 모자란 슬롯을 zero-pad합니다(반대로 모델이 더 크면 truncate). 정확한 값은 각자 check_setup.py가 찍어주는 output_features ↔ action_space를 보세요; 라이브러리·체크포인트 버전마다 드리프트할 수 있어요. 핵심은 숫자가 아니라 억지로 맞춘다는 사실입니다.

run_demo.py는 이걸 알고 일부러 처리해요: action dim이 안 맞으면 한 번 경고를 찍고 잘라내기/제로패딩 + 시뮬 범위로 clip해서 루프가 끝까지 돌게 합니다. 메커니즘은 정확히 보여주되, 정책이 이 몸에 안 맞으니 동작은 무의미 — 이 정직함이 포인트입니다.

5C.10 · Going further — 이 몸에 맞는 정책으로

base 체크포인트로 "루프를 구경"하는 게 이 폴더의 목적(=옵션 0)이에요. 여기서 제대로 동작까지 가고 싶다면, 자연스러운 다음 걸음들이 있습니다 — 대부분 5B fine-tuning과 곧장 이어져요.

5C.11 · Build & run locally

Python 3.10 권장, CPU면 충분. 순서: 설치 → check_setup.py(스펙 확인) → run_demo.py(루프).

# 스타터: robot.zip (README + requirements + check_setup.py + run_demo.py)
#   https://ksetp.netlify.app/assets/special-topics/robot.zip
unzip robot.zip && cd robot
python3.10 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt      # lerobot[smolvla] · gym-aloha · imageio[ffmpeg]

python check_setup.py                # 먼저! 스펙 출력 (첫 실행 ~1.8GB 다운로드)
python run_demo.py                   # transfer-cube, 200 steps → rollout.mp4
python run_demo.py --task insertion --steps 300
python run_demo.py --instruction "pick up the cube and pass it to the other hand"

5C.12 · The key idea — 루프가 본질, API는 곁가지

VLA = 멀티모달 LLM + action head를 닫힌 지각-행동 루프로 돌리는 것. 이 루프가 durable하고, 라이브러리 이름·키는 흐릅니다. 그리고 몸(embodiment)이 안 맞으면 정책은 안 통해요.