"AlphaSpace: (
Paper), a novel methodology designed to enhance the spatial reasoning capabilities of language models for robotic manipulation in 3D Cartesian space. AlphaSpace employs a hierarchical semantics-based tokenization strategy that encodes spatial information at both coarse and fine-grained levels. Our approach represents objects with their attributes, positions, and height information through structured tokens, enabling precise spatial reasoning without relying on traditional vision-based embeddings. This approach enables LLMs to accurately manipulate objects by positioning them at specific [x, y, z] coordinates.
1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
2import torch
3from utils import tokenize_desk, SYSTEM_PROMPT
4
5# Load the mode
6
7
8model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).to(device)
9tokenizer = AutoTokenizer.from_pretrained(model_path)
10
11# Define your workspace
12objects = [
13 {"red-cube": [51, 43, 17]},
14 {"black-cube": [44, 58, 17]},
15 {"purple-cube": [74, 59, 17]},
16 {"green-cube": [65, 82, 17]},
17]
18
19# Give a natural language instruction
20instruction = "Throw the red cube on top of the blue cylinder"
21desk, object_height = tokenize_desk(objects)
22final_instruction = SYSTEM_PROMPT.format(object_height=object_height,instruction=instruction,TABLE_MAP=desk)
23chat = [
24 {"role": "user", "content": final_instruction.strip()}
25]
26tokenized_chat = tokenizer.apply_chat_template(chat, tokenize=True, add_generation_prompt=True, use_system_prompt=False, return_tensors="pt")
27# print(len(tokenized_chat[0]))
28generated_ids = model.generate(
29 tokenized_chat.to("cuda"),
30 max_new_tokens=2048,
31 do_sample=False,
32 temperature=0.6,
33)
34# Get the solution
35result = tokenizer.decode(generated_ids[0][tokenized_chat.shape[1]:], skip_special_tokens=True)
36print(result)
We utilize
Llama-Factory library to train the model.
