Views
No views yet
1pip install einops torchvision accelerate
2pip install transformers==4.521from transformers import AutoProcessor, AutoModelForImageTextToText
2import torch
3from PIL import Image
4import requests
5from io import BytesIO
6
7ckpt = "allenai/MolmoAct-7B-D-Pretrain-0812"
8
9# load the processor
10processor = AutoProcessor.from_pretrained(
11 ckpt,
12 trust_remote_code=True,
13 torch_dtype="bfloat16",
14 device_map="auto",
15 padding_side="left",
16)
17
18# load the model
19model = AutoModelForImageTextToText.from_pretrained(
20 ckpt,
21 trust_remote_code=True,
22 torch_dtype="bfloat16",
23 device_map="auto",
24)
25
26# task instruction
27instruction = "pick orange can"
28
29# strictly follow this reasoning prompt
30prompt = (
31 f"The task is {instruction}. "
32 "What is the action that the robot should take. "
33 f"To figure out the action that the robot should take to {instruction}, "
34 "let's think through it step by step. "
35 "First, what is the depth map for this image? "
36 "Second, what is the trajectory of the end effector? "
37 "Based on the depth map of the image and the trajectory of the end effector, "
38 "what is the action that the robot should take?"
39)
40
41# apply chat template
42text = processor.apply_chat_template(
43 [
44 {
45 "role": "user",
46 "content": [dict(type="text", text=prompt)]
47 }
48 ],
49 tokenize=False,
50 add_generation_prompt=True,
51)
52
53# image observation
54url = "https://huggingface.co/allenai/MolmoAct-7B-D-Pretrain-0812/resolve/main/example.png"
55r = requests.get(url, headers={"User-Agent": "python-requests"}, timeout=30)
56r.raise_for_status()
57img = Image.open(BytesIO(r.content)).convert("RGB")
58imgs = [img]
59
60# process the image and text
61inputs = processor(
62 images=[imgs],
63 text=text,
64 padding=True,
65 return_tensors="pt",
66)
67
68# move inputs to the correct device
69inputs = {k: v.to(model.device) for k, v in inputs.items()}
70
71# generate output
72with torch.inference_mode():
73 with torch.autocast("cuda", enabled=True, dtype=torch.bfloat16):
74 generated_ids = model.generate(**inputs, max_new_tokens=256)
75
76# only get generated tokens; decode them to text
77generated_tokens = generated_ids[:, inputs['input_ids'].size(1):]
78generated_text = processor.batch_decode(generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
79
80# print the generated text
81print(f"generated text: {generated_text}")
82
83# >>> The depth map of the image is ... The trajectory of the end effector is ...
84# Based on these information, the action that the robot should take is ...
85
86# parse out all depth perception tokens
87depth = model.parse_depth(generated_text)
88print(f"generated depth perception tokens: {depth}")
89
90# >>> [ "<DEPTH_START><DEPTH_1><DEPTH_2>...<DEPTH_END>" ]
91
92# parse out all visual reasoning traces
93trace = model.parse_trace(generated_text)
94print(f"generated visual reasoning trace: {trace}")
95
96# >>> [ [[242, 115], [140, 77], [94, 58], [140, 44], [153, 26]]] ]
97
98# parse out all actions, unnormalizing with key of fractal20220817_data
99action = model.parse_action(generated_text, unnorm_key="fractal20220817_data")
100print(f"generated action: {action}")
101
102# >>> [ [0.0732076061122558, 0.08228153779226191, -0.027760173818644346,
103# 0.15932856272248652, -0.09686601126895233, 0.043916773912953344,
104# 0.996078431372549] ]1@misc{molmoact2025,
2 title={MolmoAct: Action Reasoning Models that can Reason in Space},
3 author={Jason Lee and Jiafei Duan and Haoquan Fang and Yuquan Deng and Shuo Liu and Boyang Li and Bohan Fang and Jieyu Zhang and Yi Ru Wang and Sangho Lee and Winson Han and Wilbert Pumacay and Angelica Wu and Rose Hendrix and Karen Farley and Eli VanderBilt and Ali Farhadi and Dieter Fox and Ranjay Krishna},
4 year={2025},
5 eprint={2508.07917},
6 archivePrefix={arXiv},
7 primaryClass={cs.RO},
8 url={https://arxiv.org/abs/2508.07917}
9}