왜 Cosmos 3 Edge인가?

로봇이 실제 환경에서 유연하게 동작하려면 센서, 환경, 작업에 적응할 수 있는 정책이 필요합니다. 기존의 로봇 정책은 특정 작업에만 동작하는 경우가 많았지만, 월드 모델(World Model)은 물리적 상호작용을 학습하기 위한 강력한 기반을 제공합니다. 그러나 이러한 모델은 크기가 커서 온디바이스(On-Device) 환경에서 실행하기 어려웠죠.

NVIDIA Cosmos 3 Edge는 이러한 문제를 해결합니다. 이 모델은 4B 파라미터의 옴니 모델(Omni-Model)로, Cosmos 3 제품군에 속하며, 물리적 세계 데이터로 사전 학습되었습니다. 덕분에 물체가 어떻게 움직이고 상호작용하는지에 대한 기본적인 이해를 갖추고 있습니다. 그리고 이 모델은 NVIDIA Jetson Thor에서 실행될 수 있을 만큼 작습니다.

이 튜토리얼에서는 Cosmos 3 Edge를 포스트 트레이닝(Post-Training)하여 로봇 조작 정책을 만들고, Jetson Thor에서 서빙하며, 폐루프 시뮬레이션에서 평가하는 방법을 단계별로 안내합니다.

튜토리얼 목표

  • Cosmos 3 Edge를 포스트 트레이닝하여 로봇 행동 예측
  • Jetson Thor에서 정책 서버 구동
  • 리딩 호라이즌 제어 루프에서 추론 실행
  • 폐루프 시뮬레이션에서 정책 동작 평가

모든 단계는 cosmos-framework 리포지토리에서 재현할 수 있으며, 학습된 체크포인트는 HuggingFace에서 제공됩니다.

NVIDIA Cosmos 3 Edge model running on Jetson Thor for on-device robot control Programming Illustration

포스트 트레이닝: Cosmos 3 Edge를 로봇 정책으로 만들기

데이터 준비

이 튜토리얼의 정책은 nvidia/Cosmos3-DROID 데이터셋을 사용합니다. 이 데이터셋은 약 350시간 분량의 76,000개의 성공적인 원격 조작 궤적을 포함하며, Franka Panda 팔과 Robotiq 그리퍼로 수집되었습니다.

데이터를 다운로드하고 준비하는 방법은 다음과 같습니다.

# 데이터셋 다운로드 (성공적인 시연만 포함)
hf download nvidia/Cosmos3-DROID --repo-type dataset --local-dir /path/to/Cosmos3-DROID

데이터는 3단계로 준비됩니다:

  1. 유휴 프레임 및 비작업 프레임 필터링
  2. 성공적인 시연 선택
  3. 훈련 중 랜덤 크롭, 크기 조정, 색상 지터 적용

체크포인트 변환 및 학습 실행

베이스 체크포인트를 분산 체크포인트(DCP) 형식으로 변환하고 학습을 시작합니다.

# 베이스 체크포인트를 DCP로 변환 (CPU에서 실행, GPU 불필요)
python -m cosmos_framework.scripts.convert_model_to_dcp \
  -o /path/to/Cosmos3-Edge-dcp \
  --checkpoint-path Cosmos3-Edge

# 포스트 트레이닝 실행
bash examples/launch_sft_action_policy_droid_edge.sh

Cosmos 3 Edge를 학습하려면 action_policy_droid_nano.py 설정 파일에서 다음 세 가지를 수정해야 합니다:

  • NANO_MODEL_CONFIG를 EDGE_MODEL_CONFIG로 변경
  • BASE_CHECKPOINT_PATH를 Cosmos 3 Edge DCP 체크포인트 경로로 설정
  • 스크립트 이름을 launch_sft_action_policy_droid_edge.sh로 변경

Jetson Thor에 배포

정책 서버는 WebSocket 프로토콜(OpenPI)을 통해 제공됩니다. Edge 모델은 BF16 기준 약 9GB로 Jetson Thor 메모리에 맞습니다.

export HF_TOKEN=<your_token>
# Thor에서 실행: eager 모드 사용, Triton 커널 미지원으로 dynamo 비활성화
export TORCHDYNAMO_DISABLE=1
python -m cosmos_framework.scripts.action_policy_server_robolab \
  --checkpoint_path nvidia/Cosmos3-Edge-Policy-DROID \
  --port 8000 \
  --format-prompt-as-json True

# 상태 확인
curl localhost:8000/healthz # -> OK

스모크 테스트 (로봇 없이)

import numpy as np
from PIL import Image
from openpi_client.websocket_client_policy import WebsocketClientPolicy

client = WebsocketClientPolicy(host="localhost", port=8000)

observation = {
    "prompt": "put the marker in the basket",
    "observation/wrist_image_left": np.asarray(Image.open("wrist.jpg")),
    "observation/exterior_image_1_left": np.asarray(Image.open("left.jpg")),
    "observation/exterior_image_2_left": np.asarray(Image.open("right.jpg")),
    "observation/joint_position": np.zeros(7, dtype=np.float32),  # 스모크 테스트용
    "observation/gripper_position": np.float32(0.0),  # 스모크 테스트용
}

result = client.infer(observation)
actions = result["action"]  # [32, 8]: 7개 관절 + 그리퍼, 15Hz

실제 로봇에서 실행

def get_observation():
    return {
        "prompt": "put the marker in the basket",
        "observation/wrist_image_left": read_camera("wrist"),
        "observation/exterior_image_1_left": read_camera("left"),
        "observation/exterior_image_2_left": read_camera("right"),
        "observation/joint_position": robot.joint_positions(),  # 실제 값
        "observation/gripper_position": robot.gripper_position(),
    }

while not task_done():
    result = client.infer(get_observation())
    execute_actions(result["action"][:16], hz=15)  # 32개 중 16개 실행 후 재계획

정책은 상태 조건적(State-Conditioned)이므로, 각 재계획 시 실제 관절 위치에서 시작합니다. 이는 예측 오류를 스스로 교정하는 데 도움이 됩니다.

Post-training pipeline for robot policy using world foundation model System Abstract Visual

폐루프 시뮬레이션에서 평가

실제 로봇에 적용하기 전에 시뮬레이션에서 정책을 평가하는 것이 좋습니다. RoboLab은 Isaac Lab-Arena 기반의 벤치마크로, Cosmos 3 Edge 정책 서버에 연결하여 120가지 언어 조건 조작 작업을 평가할 수 있습니다.

# RoboLab 클론 및 에셋 다운로드
git clone https://github.com/NVlabs/RoboLab.git && cd RoboLab
git lfs pull

# 시뮬레이션 이미지 빌드 및 작업 실행
./docker/build_docker.sh latest
./docker/run_docker.sh latest
python policies/cosmos3/run.py --task BananaInBowlTask

# 헤드리스로 여러 환경 실행 (성공률 통계)
python policies/cosmos3/run.py --task BananaInBowlTask --num-envs 10 --headless

폐루프 평가에서 Cosmos 3 Edge 정책은 22.9%의 성공률을 보였습니다. 이는 더 큰 Cosmos 3 Nano(36.8%)보다 낮지만, 실시간 온디바이스 추론이 가능한 점을 고려하면 실용적인 트레이드오프입니다.

이 기술의 한계 또는 주의사항

  • 계산 자원: 포스트 트레이닝은 단일 GPU 파인튜닝이 아닌 파운데이션 모델 학습입니다. 검증된 실행은 64노드의 4×GB200에서 6만 반복, 약 68시간(약 17,400 GB200시간)이 소요됩니다. 비용을 계획해야 합니다.
  • 하드웨어 요구: Jetson Thor 외에도 DGX Station과 같은 고성능 학습 환경이 필요합니다.
  • 데이터 의존성: DROID 데이터셋은 Franka Panda 로봇에 특화되어 있습니다. 다른 로봇 플랫폼을 사용한다면 자체 데이터 수집 및 설정이 필요합니다.
  • 성능 저하: Nano 대비 성공률이 낮으므로, 작업의 복잡도에 따라 더 큰 모델이 필요할 수 있습니다.

한국 개발 생태계에서의 적용 맥락

국내 로봇 공학 및 제조 현장에서는 온디바이스 AI에 대한 관심이 높아지고 있습니다. 특히, 클라우드 의존도를 줄이고 실시간 제어가 필요한 스마트 팩토리 환경에서 Cosmos 3 Edge와 같은 모델은 유망한 선택지가 될 수 있습니다. 하지만, Jetson Thor와 같은 고가의 하드웨어 도입 비용과 한국어 기반 데이터셋 부족은 해결해야 할 과제입니다. 자체 데이터를 수집할 때는 Python 기반 데이터 처리 파이프라인을 참고하여 효율적으로 관리하는 것이 좋습니다.

Closed-loop simulation evaluation of Cosmos 3 Edge policy in RoboLab Coding Session Visual

결론: 실무 적용 조언

Cosmos 3 Edge는 온디바이스 로봇 제어를 위한 실용적인 월드 모델입니다. 포스트 트레이닝을 통해 특정 로봇 작업에 맞게 조정할 수 있으며, Jetson Thor에서 실시간으로 실행됩니다. 폐루프 시뮬레이션을 통해 안전하게 성능을 검증할 수 있습니다.

함께 보면 좋은 글

다음 단계 학습 방향

  • Cosmos 3 Edge 모델 카드를 확인하여 지원하는 다양한 로봇 플랫폼 탐색
  • RoboLab의 120가지 작업을 모두 실행하여 정책의 강점/약점 파악
  • 자체 로봇 데이터 수집 파이프라인 구축 및 LeRobot Dataset v3 변환 방법 학습
  • 원문 튜토리얼에서 최신 정보 확인
본 콘텐츠는 신뢰할 수 있는 출처를 바탕으로 AI 도구를 활용하여 초안이 작성되었으며, 편집자의 검토를 거쳐 발행되었습니다. 전문가의 조언을 대체하지 않습니다.