utils/constants.py

"""Centralized constants: task metadata and robot configuration."""

# ========== Task Metadata ==========
# task_name -> prompt, robot_type mapping
TASK_METADATA = {
    # ac_one (arx5)
    "arrange_flowers": {
        "prompt": "Put the 4 flowers into the vase.",
        "robot_type": "arx5",
    },
    "hang_the_cup": {
        "prompt": "Hang the cup on the rack.",
        "robot_type": "arx5",
    },
    "pick_out_the_green_blocks": {
        "prompt": "Find all the green blocks and put them into the basket.",
        "robot_type": "arx5",
    },
    "press_the_button": {
        "prompt": "Press the buttons in the following sequence: pink, blue, green, and then yellow.",
        "robot_type": "arx5",
    },
    "turn_on_the_light_switch": {
        "prompt": "Turn on the lamp.",
        "robot_type": "arx5",
    },
    "water_the_flowers": {
        "prompt": "Water the potted plants.",
        "robot_type": "arx5",
    },
    "wipe_the_table": {
        "prompt": "Wipe the stains off the desk with a rag.",
        "robot_type": "arx5",
    },
    # ur5
    "arrange_fruits": {
        "prompt": "Arrange the fruit in the basket.",
        "robot_type": "ur5",
    },
    "item_classification": {
        "prompt": "Place the stationery in the yellow box and the electronics in the blue box.",
        "robot_type": "ur5",
    },
    "shred_paper": {
        "prompt": "Put the paper into the shredder.",
        "robot_type": "ur5",
    },
    # w1
    "fold_the_clothes": {
        "prompt": "Fold the T-shirts and stack them neatly in the upper-left corner of the table.",
        "robot_type": "w1_endpose",
    },
    "hold_the_tray_with_both_hands": {
        "prompt": "Place the ball from the desk onto the small tray, and then move the small tray onto the large tray.",
        "robot_type": "w1_endpose",
    },
    "place_objects_into_desk_drawer": {
        "prompt": "Open the drawer, put the bottle opener inside, and close the drawer.",
        "robot_type": "w1_endpose",
    },
    "put_in_pen_container": {
        "prompt": "Put the pens on the desk into the pen holder.",
        "robot_type": "w1_endpose",
    },
    "put_the_shoes_back": {
        "prompt": "Pair the two pairs of shoes on the desk and place them on the shoe rack",
        "robot_type": "w1_endpose",
    },
    "stack_bowls": {
        "prompt": "Put the blue bowl into the beige bowl, and put the green bowl into the blue bowl.",
        "robot_type": "w1_endpose",
    },
    "sweep_the_trash": {
        "prompt": "Sweep the trash on the table into the dustpan.",
        "robot_type": "w1_endpose",
    },
    "tidy_up_the_makeup_table": {
        "prompt": "Sort and organize the cosmetics on the table.",
        "robot_type": "w1_endpose",
    },
    "tie_a_knot": {
        "prompt": "Tie a knot with the string on the table.",
        "robot_type": "w1_endpose",
    },
    "untie_the_shoelaces": {
        "prompt": "Remove the laces from the shoes, then place them on the table.",
        "robot_type": "w1_endpose",
    },
    # aloha
    "lint_roller_remove_dirt": {
        "prompt": "Use a lint roller to remove the debris from the clothing.",
        "robot_type": "aloha_joint",
    },
    "pack_the_items": {
        "prompt": "Box up the tablet and its accessories.",
        "robot_type": "aloha_joint",
    },
    "pack_the_toothbrush_holder": {
        "prompt": "Put the toothbrush and toothpaste into the toiletries case in sequence, close the case, and then place it into the basket.",
        "robot_type": "aloha_joint",
    },
    "paint_jam": {
        "prompt": "Spread the bread with jam.",
        "robot_type": "aloha_joint",
    },
    "put_the_books_back": {
        "prompt": "Place the books back onto the bookshelf.",
        "robot_type": "aloha_joint",
    },
    "put_the_pencil_case_into_the_schoolbag": {
        "prompt": "Put the pencil case into the backpack.",
        "robot_type": "aloha_joint",
    },
    "scoop_with_a_small_spoon": {
        "prompt": "Scoop beans and place them into the empty bowl.",
        "robot_type": "aloha_joint",
    },
    "stamp_positioning": {
        "prompt": "Stamp the signature area on the paper.",
        "robot_type": "aloha_endpose",
    },
    "wipe_the_blackboard": {
        "prompt": "Wipe the blackboard clean.",
        "robot_type": "aloha_joint",
    },
    "wrap_with_a_soft_cloth": {
        "prompt": "Bundle the objects together using the cloth on the table.",
        "robot_type": "aloha_joint",
    },
}

# ========== Robot Configuration ==========
# robot_type -> image_type list and action_type
#
# image_type uses the raw camera key names from the robot API (README#get-state).
# cam_key_map defines how each raw camera key maps to the logical slot
# used by _build_frame: "high", "left", "right".
#   aloha / w1 : cam_high -> high, cam_left_wrist -> left, cam_right_wrist -> right
#   arx5       : cam_global -> high, cam_arm -> left, cam_side -> right
#   ur5        : cam_global -> high, cam_arm -> left, cam_arm (copy) -> right
ROBOT_CONFIGS = {
    "aloha_joint": {
        "image_type": ["cam_left_wrist", "cam_right_wrist", "cam_high"],
        "cam_key_map": {
            "cam_high": "high",
            "cam_left_wrist": "left",
            "cam_right_wrist": "right",
        },
        "action_type": "joint",     # dual-arm: left+right joints concatenated
        "num_arms": 2,              # joint space, no eef inverse needed
    },
    "aloha_endpose": {
        "image_type": ["cam_left_wrist", "cam_right_wrist", "cam_high"],
        "cam_key_map": {
            "cam_high": "high",
            "cam_left_wrist": "left",
            "cam_right_wrist": "right",
        },
        "action_type": "pos",     # dual-arm: left+right joints concatenated
        "num_arms": 2,              # joint space, no eef inverse needed
    },
    "arx5": {
        "image_type": ["cam_global", "cam_arm", "cam_side"],
        "cam_key_map": {
            "cam_global": "high",
            "cam_arm": "left",
            "cam_side": "right",
        },
        "action_type": "leftpos",   # single-arm
        "num_arms": 1,
    },
    "ur5": {
        "image_type": ["cam_global", "cam_arm"],
        "cam_key_map": {
            "cam_global": "high",
            "cam_arm": "left",
            # cam_arm is also copied to "right" in _build_frame for ur5
        },
        "action_type": "leftpos",   # single-arm
        "num_arms": 1,
    },
    "w1_joint": {
        "image_type": ["cam_left_wrist", "cam_right_wrist", "cam_high"],
        "cam_key_map": {
            "cam_high": "high",
            "cam_left_wrist": "left",
            "cam_right_wrist": "right",
        },
        "action_type": "joint",       # dual-arm: left+right joints concatenated
        "num_arms": 2,
    },
    "w1_endpose": {
        "image_type": ["cam_left_wrist", "cam_right_wrist", "cam_high"],
        "cam_key_map": {
            "cam_high": "high",
            "cam_left_wrist": "left",
            "cam_right_wrist": "right",
        },
        "action_type": "pos",       # dual-arm: left+right joints concatenated
        "num_arms": 2,
    },
}

SUPPORTED_ROBOTS = list(ROBOT_CONFIGS.keys())