| """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") |
|
|
| |
| 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: |
| 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) |
|
|
