File size: 3,416 Bytes
7399b6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """LLM-based dual-arm task allocation (the 'T' in TAMP: decide WHICH arm does WHAT).
Given the scene (object world positions + the two YAM arm bases) and a goal, ask an LLM
to output a per-arm plan: which arm grasps which object, in what order, and whether a
hand-off is needed. Reachability heuristic + collision avoidance are described to the model.
The motion for each assigned sub-task is then executed by the existing PRM + diff-IK layer.
Uses HuggingFace router (OpenAI-compatible). Key via env HF_TOKEN. No key is printed.
"""
import os, json, urllib.request
HF = os.environ["HF_TOKEN"]
MODEL = os.environ.get("VLM_MODEL", "deepseek-ai/DeepSeek-V4-Flash")
# ---- real scene (env-local world frame; z up; table top ~0.45) ----
SCENE = {
"arms": {
"left_robot": {"base": [-0.2, 0.2, 0.45], "reach_m": 0.55},
"right_robot": {"base": [-0.2, -0.2, 0.45], "reach_m": 0.55},
},
"objects": {
"apple": {"pos": [0.00, 0.00, 0.47], "graspable": "sphere ~4cm"},
"grape": {"pos": [0.00, 0.10, 0.44], "graspable": "cluster ~15x8cm"},
"can": {"pos": [0.00, -0.12, 0.50], "graspable": "cylinder d4.4 h11cm"},
},
"placement_targets": {"cart": [-0.5, 0.0, 0.45]},
"notes": "Both arms are fixed-base. An arm reaches an object if |object - base| < reach_m. "
"Prefer assigning each object to the CLOSER arm; the two arms must not be sent to "
"the same object at the same time (collision). Grasp is top-down.",
}
GOAL = os.environ.get("GOAL", "Pick up all three objects and place them into the cart, using both arms in parallel where possible.")
SYS = (
"You are the task-allocation planner for a BIMANUAL robot (two 6-DoF arms, left_robot and "
"right_robot, each with a parallel-jaw gripper). Decide which arm performs which sub-task. "
"Respect reachability (|obj-base|<reach) and avoid sending both arms to the same object at once. "
"Return STRICT JSON only, no prose, with schema:\n"
"{\"assignments\":[{\"arm\":\"left_robot|right_robot\",\"object\":\"...\",\"action\":\"pick_and_place\","
"\"target\":\"cart\",\"order\":1,\"reason\":\"...\"}],\"parallelizable\":[[...],[...]],"
"\"handoffs\":[],\"notes\":\"...\"}"
)
USER = f"SCENE:\n{json.dumps(SCENE, indent=2)}\n\nGOAL: {GOAL}\n\nReturn the JSON plan."
body = json.dumps({
"model": MODEL,
"messages": [{"role": "system", "content": SYS}, {"role": "user", "content": USER}],
"temperature": 0.2, "max_tokens": 900,
}).encode()
req = urllib.request.Request("https://router.huggingface.co/v1/chat/completions", data=body,
headers={"Authorization": f"Bearer {HF}", "Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=90) as r:
out = json.load(r)
msg = out["choices"][0]["message"]["content"]
print("=== RAW LLM OUTPUT ===")
print(msg)
# try to extract JSON
try:
s = msg[msg.index("{"): msg.rindex("}")+1]
plan = json.loads(s)
print("\n=== PARSED PLAN ===")
for a in plan.get("assignments", []):
print(f" [{a.get('order')}] {a.get('arm')}: {a.get('action')} {a.get('object')} -> {a.get('target')} ({a.get('reason','')[:70]})")
print(" parallelizable:", plan.get("parallelizable"))
json.dump(plan, open("outputs/dualarm_allocation.json", "w"), indent=2)
print("\nsaved -> outputs/dualarm_allocation.json")
except Exception as e:
print("\n[parse warn]", e)
|