AffordanceNet/phaseaware_data_construction/test_instruction_parser.py

"""Unit tests for instruction_parser: cover all 40 LIBERO task instructions."""

import pytest
from phaseaware_data_construction.instruction_parser import Subgoal, parse_instruction


# ---------- helpers ----------
def types(sgs):
    return [s.type for s in sgs]


def objects(sgs):
    return [s.object for s in sgs]


# ==================== libero_object (10 tasks) ====================
# All: "pick up X and place it in the basket" → [pick, place]

LIBERO_OBJECT_INSTRUCTIONS = [
    "pick up the orange juice and place it in the basket",
    "pick up the ketchup and place it in the basket",
    "pick up the cream cheese and place it in the basket",
    "pick up the bbq sauce and place it in the basket",
    "pick up the alphabet soup and place it in the basket",
    "pick up the milk and place it in the basket",
    "pick up the salad dressing and place it in the basket",
    "pick up the butter and place it in the basket",
    "pick up the tomato sauce and place it in the basket",
    "pick up the chocolate pudding and place it in the basket",
]

LIBERO_OBJECT_EXPECTED_OBJECTS = [
    "orange juice", "ketchup", "cream cheese", "bbq sauce",
    "alphabet soup", "milk", "salad dressing", "butter",
    "tomato sauce", "chocolate pudding",
]


@pytest.mark.parametrize("idx", range(10))
def test_libero_object(idx):
    sgs = parse_instruction(LIBERO_OBJECT_INSTRUCTIONS[idx])
    assert types(sgs) == ["pick", "place"]
    assert sgs[0].object == LIBERO_OBJECT_EXPECTED_OBJECTS[idx]
    assert sgs[1].object == LIBERO_OBJECT_EXPECTED_OBJECTS[idx]
    assert sgs[1].target == "basket"
    # pick has no spatial descriptor; place uses "in" relation
    assert sgs[0].descriptor == ""
    assert sgs[1].relation == "in"


# ==================== libero_spatial (10 tasks) ====================
# All: "pick up the black bowl {spatial} and place it on the plate" → [pick, place]

LIBERO_SPATIAL_INSTRUCTIONS = [
    "pick up the black bowl next to the cookie box and place it on the plate",
    "pick up the black bowl in the top drawer of the wooden cabinet and place it on the plate",
    "pick up the black bowl on the ramekin and place it on the plate",
    "pick up the black bowl on the stove and place it on the plate",
    "pick up the black bowl between the plate and the ramekin and place it on the plate",
    "pick up the black bowl on the cookie box and place it on the plate",
    "pick up the black bowl next to the plate and place it on the plate",
    "pick up the black bowl next to the ramekin and place it on the plate",
    "pick up the black bowl from table center and place it on the plate",
    "pick up the black bowl on the wooden cabinet and place it on the plate",
]

LIBERO_SPATIAL_EXPECTED_DESCRIPTORS = [
    "next to cookie box",
    "in top drawer of wooden cabinet",
    "on ramekin",
    "on stove",
    "between plate and ramekin",
    "on cookie box",
    "next to plate",
    "next to ramekin",
    "from table center",
    "on wooden cabinet",
]


@pytest.mark.parametrize("idx", range(10))
def test_libero_spatial(idx):
    sgs = parse_instruction(LIBERO_SPATIAL_INSTRUCTIONS[idx])
    assert types(sgs) == ["pick", "place"]
    assert sgs[0].object == "black bowl"
    assert sgs[1].object == "black bowl"
    assert sgs[1].target == "plate"
    # pick carries spatial descriptor; place uses "on" relation
    assert sgs[0].descriptor == LIBERO_SPATIAL_EXPECTED_DESCRIPTORS[idx]
    assert sgs[1].relation == "on"


# ==================== libero_goal (10 tasks) ====================

class TestLiberoGoal:
    def test_put_bowl_on_plate(self):
        sgs = parse_instruction("put the bowl on the plate")
        assert types(sgs) == ["pick", "place"]
        assert sgs[0].object == "bowl"
        assert sgs[1].target == "plate"
        assert sgs[1].relation == "on"

    def test_put_wine_bottle_on_rack(self):
        sgs = parse_instruction("put the wine bottle on the rack")
        assert types(sgs) == ["pick", "place"]
        assert sgs[0].object == "wine bottle"
        assert sgs[1].target == "rack"

    def test_open_drawer_put_inside(self):
        """open X and put Y inside → [open, pick, place] with target=X"""
        sgs = parse_instruction("open the top drawer and put the bowl inside")
        assert types(sgs) == ["open", "pick", "place"]
        assert sgs[0].object == "top drawer"
        assert sgs[1].object == "bowl"
        assert sgs[2].object == "bowl"
        assert sgs[2].target == "top drawer"  # "inside" resolves to opened drawer

    def test_put_cream_cheese_in_bowl(self):
        sgs = parse_instruction("put the cream cheese in the bowl")
        assert types(sgs) == ["pick", "place"]
        assert sgs[0].object == "cream cheese"
        assert sgs[1].target == "bowl"

    def test_put_on_top_of(self):
        sgs = parse_instruction("put the wine bottle on top of the cabinet")
        assert types(sgs) == ["pick", "place"]
        assert sgs[0].object == "wine bottle"
        assert sgs[1].target == "cabinet"
        assert sgs[1].relation == "on top of"

    def test_push(self):
        sgs = parse_instruction("push the plate to the front of the stove")
        assert types(sgs) == ["push"]
        assert sgs[0].object == "plate"
        assert sgs[0].relation == "to the front of"

    def test_turn_on(self):
        sgs = parse_instruction("turn on the stove")
        assert types(sgs) == ["turn_on"]
        assert sgs[0].object == "stove"

    def test_put_bowl_on_stove(self):
        sgs = parse_instruction("put the bowl on the stove")
        assert types(sgs) == ["pick", "place"]
        assert sgs[1].target == "stove"

    def test_put_bowl_on_top_of_cabinet(self):
        sgs = parse_instruction("put the bowl on top of the cabinet")
        assert types(sgs) == ["pick", "place"]
        assert sgs[1].target == "cabinet"

    def test_open_middle_drawer(self):
        sgs = parse_instruction("open the middle drawer of the cabinet")
        assert types(sgs) == ["open"]
        assert "middle drawer" in sgs[0].object


# ==================== libero_10 (10 tasks) ====================

class TestLibero10:
    def test_put_two_mugs(self):
        """put X on A and put Z on B → 4 subgoals"""
        sgs = parse_instruction(
            "put the white mug on the left plate and put the yellow and white mug on the right plate"
        )
        assert types(sgs) == ["pick", "place", "pick", "place"]
        assert sgs[0].object == "white mug"
        assert sgs[1].target == "left plate"
        assert sgs[2].object == "yellow and white mug"
        assert sgs[3].target == "right plate"

    def test_put_mug_and_pudding(self):
        sgs = parse_instruction(
            "put the white mug on the plate and put the chocolate pudding to the right of the plate"
        )
        assert types(sgs) == ["pick", "place", "pick", "place"]
        assert sgs[2].object == "chocolate pudding"
        assert sgs[1].relation == "on"
        assert sgs[3].relation == "to the right of"

    def test_put_and_close(self):
        """put X in Y and close it → [pick, place, close_Y]"""
        sgs = parse_instruction(
            "put the yellow and white mug in the microwave and close it"
        )
        assert types(sgs) == ["pick", "place", "close"]
        assert sgs[0].object == "yellow and white mug"
        assert sgs[1].target == "microwave"
        assert sgs[1].relation == "in"
        assert sgs[2].object == "microwave"  # "it" → last_target

    def test_turn_on_and_put_on_it(self):
        """turn on X and put Y on it → [turn_on, pick, place on X]"""
        sgs = parse_instruction("turn on the stove and put the moka pot on it")
        assert types(sgs) == ["turn_on", "pick", "place"]
        assert sgs[0].object == "stove"
        assert sgs[1].object == "moka pot"
        assert sgs[2].target == "stove"  # "it" → last_object (turn_on has no target)

    def test_put_both_with_and(self):
        """put both X and Z in Y → 4 subgoals"""
        sgs = parse_instruction(
            "put both the alphabet soup and the cream cheese box in the basket"
        )
        assert types(sgs) == ["pick", "place", "pick", "place"]
        assert sgs[0].object == "alphabet soup"
        assert sgs[2].object == "cream cheese box"
        assert sgs[1].target == "basket"
        assert sgs[3].target == "basket"
        assert sgs[1].relation == "in"
        assert sgs[3].relation == "in"

    def test_put_both_tomato(self):
        sgs = parse_instruction(
            "put both the alphabet soup and the tomato sauce in the basket"
        )
        assert types(sgs) == ["pick", "place", "pick", "place"]
        assert sgs[0].object == "alphabet soup"
        assert sgs[2].object == "tomato sauce"

    def test_put_both_plural_no_and(self):
        """put both [plural noun] → 4 subgoals with duplicated object"""
        sgs = parse_instruction("put both moka pots on the stove")
        assert types(sgs) == ["pick", "place", "pick", "place"]
        assert sgs[0].object == "moka pot"
        assert sgs[2].object == "moka pot"
        assert sgs[1].target == "stove"

    def test_put_both_butter(self):
        sgs = parse_instruction(
            "put both the cream cheese box and the butter in the basket"
        )
        assert types(sgs) == ["pick", "place", "pick", "place"]
        assert sgs[0].object == "cream cheese box"
        assert sgs[2].object == "butter"

    def test_put_bowl_close_drawer(self):
        sgs = parse_instruction(
            "put the black bowl in the bottom drawer of the cabinet and close it"
        )
        assert types(sgs) == ["pick", "place", "close"]
        assert sgs[0].object == "black bowl"
        assert "bottom drawer" in sgs[1].target
        assert sgs[1].relation == "in"
        assert "bottom drawer" in sgs[2].object  # "it" → last_target

    def test_pick_place_caddy(self):
        sgs = parse_instruction(
            "pick up the book and place it in the back compartment of the caddy"
        )
        assert types(sgs) == ["pick", "place"]
        assert sgs[0].object == "book"
        assert sgs[1].object == "book"
        assert "back compartment" in sgs[1].target
        assert sgs[0].descriptor == ""
        assert sgs[1].relation == "in"