공간 인식과 로봇 행동 예측을 위한 MolmoAct 구현 튜토리얼

Artificial Intelligence Applications Technology Computer Vision Editors Pick Machine Learning Physical AI Staff Tutorials In this tutorial, we walk through MolmoAct step by step and build a practical understanding of how action-reasoning models can reason in space from visual observations. We set up the environment, load the model, prepare multi-view image inputs, and explore how MolmoAct produces depth-aware reasoning, visual traces, and actionable robot outputs from natural language instructions. As we move through the workflow, we run inference and also examine how the model parses actions, visualizes trajectories, and supports more advanced processing pipelines for robotics-oriented tasks. Copy Code Copied Use a different Browser print("=" * 80) print("🔧 SECTION 1: INSTALLATION AND SETUP") print("=" * 80) import subprocess import sys def install_packages(): """Install all required packages for MolmoAct""" packages = [ "torch>=2.0.0", "torchvision", "transformers==4.52", "accelerate", "einops", "Pillow", "numpy", "matplotlib", "requests", "scipy", "huggingface_hub", ] for package in packages: print(f"📦 Installing {package}...") subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package]) print("✅ All packages installed successfully!") install_packages() print("\n" + "=" * 80) print("📚 SECTION 2: IMPORTS AND CONFIGURATION") print("=" * 80) import torch import numpy as np import matplotlib.pyplot as plt from PIL import Image import requests from io import BytesIO from typing import List, Tuple, Dict, Optional, Union import json import time from dataclasses import dataclass import warnings import re warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=UserWarning) print(f"🖥️ Device: {device}") if torch.cuda.is_available(): print(f"🎮 GPU: {torch.cuda.get_device_name(0)}") print(f"💾 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB") print("\n" + "=" * 80) print("🤖 SECTION 3: MOLMOACT MODEL LOADER") print("=" * 80) @dataclass class MolmoActConfig: """Configuration for MolmoAct model""" model_name: str = "allenai/MolmoAct-7B-D-0812" torch_dtype: str = "bfloat16" device_map: str = "auto" max_new_tokens: int = 256 temperature: float = 0.0 do_sample: bool = False We set up the tutorial and prepared the environment needed to run MolmoAct in Google Colab. We install all required packages, import the core libraries, and configure the runtime to detect whether GPU acceleration is available. We also define the base configuration class that stores the main model settings we use throughout the rest of the tutorial. Copy Code Copied Use a different Browser class MolmoActModel: """ MolmoAct Model Wrapper for Easy Inference This class provides a high-level interface for: - Loading and managing the model - Running inference with proper prompting - Parsing outputs (depth, trace, actions) - Batch processing """ def __init__(self, config: Optional[MolmoActConfig] = None): self.config = config or MolmoActConfig() self.model = None self.processor = None self._loaded = False def load(self) -> None: """Load the MolmoAct model and processor""" if self._loaded: print("⚠️ Model already loaded!") return print(f"🔄 Loading MolmoAct model: {self.config.model_name}") print(" This may take a few minutes on first run...") from transformers import AutoModelForImageTextToText, AutoProcessor dtype = getattr(torch, self.config.torch_dtype) print(" 📥 Loading model weights...") self.model = AutoModelForImageTextToText.from_pretrained( self.config.model_name, trust_remote_code=True, torch_dtype=dtype, device_map=self.config.device_map, ) print(" 📥 Loading processor...") try: self.processor = AutoProcessor.from_pretrained( self.config.model_name, trust_remote_code=True, ) if hasattr(self.processor, 'tokenizer'): self.processor.tokenizer.padding_side = "left" except TypeError as e: if "prompt_templates" in str(e): print(" ⚠️ Handling custom processor configuration...") from transformers.dynamic_module_utils import get_class_from_dynamic_module processor_class = get_class_from_dynamic_module( "processing_molmoact.MolmoActProcessor", self.config.model_name, trust_remote_code=True, ) from transformers import AutoTokenizer, AutoImageProcessor tokenizer = AutoTokenizer.from_pretrained( self.config.model_name, trust_remote_code=True, padding_side="left", ) image_processor = AutoImageProcessor.from_pretrained( self.config.model_name, trust_remote_code=True, ) self.processor = processor_class( image_processor=image_processor, tokenizer=tokenizer, ) else: raise e self._loaded = True print("✅ Model loaded successfully!") self._print_model_info() def _print_model_info(self) -> None: """Print model information""" total_params = sum(p.numel() for p in self.model.parameters()) trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) print(f"\n📊 Model Statistics:") print(f" Total Parameters: {total_params / 1e9:.2f}B") print(f" Trainable Parameters: {trainable_params / 1e9:.2f}B") print(f" Model dtype: {next(self.model.parameters()).dtype}") def build_prompt(self, instruction: str) -> str: """ Build the reasoning prompt for MolmoAct The prompt structure is crucial for MolmoAct to generate: 1. Depth perception tokens 2. Visual trajectory trace 3. Action predictions """ prompt = ( f"The task is {instruction}. " "What is the action that the robot should take. " f"To figure out the action that the robot should take to {instruction}, " "let's think through it step by step. " "First, what is the depth map for the first image? " "Second, what is the trajectory of the end effector in the first image? " "Based on the depth map of the first image and the trajectory of the end effector in the first image, " "along with other images from different camera views as additional information, " "what is the action that the robot should take?" ) return prompt We begin building the main MolmoAct model wrapper that makes inference easier to manage. We load the model and processor, handle custom processor initialization logic, and print useful model statistics once loading is complete. We also define a prompt-building method that helps us structure the reasoning query to guide the model toward depth, trace, and action generation. Copy Code Copied Use a different Browser @torch.inference_mode() def generate( self, images: List[Image.Image], instruction: str, max_new_tokens: Optional[int] = None, ) -> Dict: """ Generate action reasoning from images and instruction Args: images: List of PIL Images instruction: Task instruction max_new_tokens: Override default max tokens Returns: Dictionary containing: - text: Generated reasoning text - depth: Parsed depth tokens - trace: Parsed visual trace coordinates - action: Parsed action values """ if not self._loaded: raise RuntimeError("Model not loaded! Call .load() first.") prompt = self.build_prompt(instruction) max_tokens = max_new_tokens or self.config.max_new_tokens text = self.processor.apply_chat_template( [{"role": "user", "content": [dict(type="text", text=prompt)]}], tokenize=False, add_generation_prompt=True, ) inputs = self.processor( images=[images], text=text, padding=True, return_tensors="pt", ) inputs = {k: v.to(self.model.device) for k, v in inputs.items()} with torch.autocast("cuda", enabled=True, dtype=torch.bfloat16): generated_ids = self.model.generate( **inputs, max_new_tokens=max_tokens, do_sample=self.config.do_sample, ) generated_tokens = generated_ids[:, inputs['input_ids'].size(1):] generated_text = self.processor.batch_decode( generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False )[0] result = { "text": generated_text, "depth": self._safe_parse_depth(generated_text), "trace": self._safe_parse_trace(generated_text), "action": self._safe_parse_action(generated_text, unnorm_key="molmoact"), "action_raw": self._safe_parse_action(generated_text, unnorm_key=None), } return result def _safe_parse_depth(self, text: