[실습] 정밀 모방학습 — ACT·Diffusion Policy를 LeRobot으로

들어가며

3편에서 다룬 ACT와 Diffusion Policy는 모두 LeRobot(HuggingFace의 로봇 학습 라이브러리)에 사전학습 모델이 통합되어 있습니다. 즉 로봇 한 대도, GPU도 없이 노트북의 시뮬레이션으로 굴려볼 수 있습니다.

이번 글에서 직접 굴려볼 두 가지는 다음과 같습니다.

  1. Diffusion Policy로 PushT 풀기 — 가장 단순한 2D 푸시 작업. 5분이면 첫 결과를 봅니다.
  2. ACT로 ALOHA 시뮬레이션 양팔 작업 — 큐브를 한 팔에서 다른 팔로 넘기기

GPU가 있으면 추론이 빠르지만, CPU만으로도 가능합니다. 한 episode + 영상 저장이 1~3분 걸립니다.

0. 환경 준비

LeRobot은 빠르게 자라고 있어 git clone + 편집 가능 설치가 가장 안전합니다.

# LeRobot 클론과 설치
git clone https://github.com/huggingface/lerobot.git
cd lerobot
pip install -e ".[pusht,aloha]"

설치 확인:

import lerobot
print(lerobot.__version__)

pip install lerobot 한 줄로도 설치 가능하나, 자주 업데이트되므로 git clone을 권장합니다.

1. Diffusion Policy로 PushT 풀기

PushT는 LeRobot의 시작점입니다. 화면 위 T자 도형을 마커 위치로 밀어 넣는 단순한 2D 작업이지만, multimodal 행동의 본질을 보여줍니다.

import torch
import gym_pusht  # PushT 환경 등록
import gymnasium as gym
import imageio

from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy

# 사전학습 모델 다운로드 (~50MB)
device = "cuda" if torch.cuda.is_available() else "cpu"
policy = DiffusionPolicy.from_pretrained("lerobot/diffusion_pusht")
policy.to(device).eval()

env = gym.make(
    "gym_pusht/PushT-v0",
    obs_type="pixels_agent_pos",
    max_episode_steps=300,
)

obs, info = env.reset(seed=42)
frames = [env.render()]
done = False
total_reward = 0

while not done:
    state = torch.from_numpy(obs["agent_pos"]).float().to(device).unsqueeze(0)
    image = (
        torch.from_numpy(obs["pixels"]).float().permute(2, 0, 1).unsqueeze(0).to(device) / 255.0
    )
    observation = {
        "observation.state": state,
        "observation.image": image,
    }
    with torch.no_grad():
        action = policy.select_action(observation)

    obs, reward, terminated, truncated, info = env.step(action.squeeze(0).cpu().numpy())
    done = terminated or truncated
    frames.append(env.render())
    total_reward += reward

print(f"total_reward: {total_reward:.2f}, success: {info.get('is_success', False)}")
imageio.mimsave("pusht_diffusion.mp4", frames, fps=30)

결과 해석

pusht_diffusion.mp4를 열어보면, 작은 원이 T자 도형을 미는 모습이 보입니다. 잘 학습된 Diffusion Policy는 보통 몇 번의 시도 안에 T를 목표 위치에 안착시킵니다.

total_reward는 0(완전 실패)~1(완벽 성공) 사이입니다. 이 사전학습 모델은 평균 0.85+ 정도가 정상입니다.

같은 코드를 다른 seed로 여러 번 돌려보시기 바랍니다. Diffusion Policy의 multimodal 특성이 보입니다 — 같은 시작 상태에서도 매번 조금씩 다른 경로로 T를 밉니다. 평균만 학습한 BC라면 항상 같은 길로 갔을 겁니다.

2. ACT로 ALOHA 시뮬레이션

ALOHA 환경은 양팔 작업입니다. LeRobot에는 cube transfer(한 팔에서 다른 팔로 큐브 넘기기) 같은 task가 사전학습된 ACT 모델과 함께 들어 있습니다.

import torch
import gym_aloha  # ALOHA 환경 등록
import gymnasium as gym
import imageio

from lerobot.common.policies.act.modeling_act import ACTPolicy

device = "cuda" if torch.cuda.is_available() else "cpu"

policy = ACTPolicy.from_pretrained("lerobot/act_aloha_sim_transfer_cube_human")
policy.to(device).eval()

env = gym.make(
    "gym_aloha/AlohaTransferCube-v0",
    obs_type="pixels_agent_pos",
    max_episode_steps=400,
)

obs, info = env.reset(seed=42)
frames = [env.render()]
done = False
total_reward = 0

while not done:
    state = torch.from_numpy(obs["agent_pos"]).float().to(device).unsqueeze(0)
    images = {
        f"observation.images.{cam}":
            torch.from_numpy(obs["pixels"][cam])
                .float().permute(2, 0, 1).unsqueeze(0).to(device) / 255.0
        for cam in obs["pixels"]
    }
    observation = {"observation.state": state, **images}

    with torch.no_grad():
        action = policy.select_action(observation)

    obs, reward, terminated, truncated, info = env.step(action.squeeze(0).cpu().numpy())
    done = terminated or truncated
    frames.append(env.render())
    total_reward += reward

print(f"reward: {total_reward:.2f}, success: {info.get('is_success', False)}")
imageio.mimsave("aloha_act.mp4", frames, fps=50)

결과 해석

영상에는 양팔이 큐브를 들어 한 팔에서 다른 팔로 넘기는 모습이 나옵니다. ACT는 3편에서 본 것처럼 다음 100 step의 행동을 한꺼번에 예측하는 action chunking을 사용합니다.

policy.select_action 안에서 chunk가 100개 행동을 출력해도 그 중 첫 행동만 반환됩니다. 다음 step에서는 새 관측으로 다시 chunk 예측 → 첫 행동 반환을 반복하는 구조입니다. ACT의 temporal ensembling이 자동으로 적용됩니다.

3. 두 모델 비교 — multimodality 직접 보기

같은 환경에서 다른 seed로 여러 episode를 돌리고 마지막 경로만 그려보면 multimodal action distribution의 의미가 명확해집니다.

import matplotlib.pyplot as plt

paths = []
for seed in [0, 1, 2, 3, 4]:
    obs, _ = env.reset(seed=seed)
    coords = []
    for _ in range(300):
        # ... 위와 같은 추론 루프 ...
        coords.append(obs["agent_pos"].copy())
    paths.append(coords)

for i, p in enumerate(paths):
    p = np.array(p)
    plt.plot(p[:, 0], p[:, 1], label=f"seed {i}")
plt.legend()
plt.savefig("paths.png")

Diffusion Policy의 경로는 매번 다르고, 각각이 자연스럽습니다. 평균만 학습한 BC라면 보통 비슷한 경로만 반복합니다.

정리

실습손에 잡힌 개념
Diffusion Policy + PushT행동 분포를 생성 모델로 학습한다는 의미
ACT + ALOHAaction chunking이 실제로 어떻게 적용되는지
경로 비교multimodal vs unimodal의 시각적 차이

Diffusion Policy와 ACT는 5편의 OpenVLA와 6편의 π₀의 직접적 조상입니다. 다음 실습에서는 OpenVLA 같은 거대 VLA를 같은 방식으로 굴려봅니다.

다음 글 안내

LinkedIn