AffordanceNet/phaseaware_data_construction/instruction_parser.py

"""Instruction parser: extract ordered subgoal sequences from LIBERO task instructions."""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import List, Optional


@dataclass
class Subgoal:
    type: str  # "pick", "place", "open", "close", "turn_on", "push"
    object: str
    target: str = ""
    descriptor: str = ""   # Spatial qualifier for object (e.g. "next to cookie box")
    relation: str = ""     # Preposition for placement (e.g. "to the right of")

    def label(self) -> str:
        """Label for phase matching. Does NOT include descriptor/relation."""
        parts = [self.type, self.object]
        if self.target:
            parts.append(self.target)
        return "_".join(parts).replace(" ", "_")


# Verb patterns ordered by specificity (multi-word first)
VERB_PATTERNS = [
    ("pick up", "pick_up"),
    ("turn on", "turn_on"),
    ("turn off", "turn_off"),
    ("put", "put"),
    ("place", "place"),
    ("open", "open"),
    ("close", "close"),
    ("push", "push"),
    ("stack", "stack"),
]

# Prepositions that separate object from target
# Multi-word patterns must come before single-word ones to match greedily
PREP_PATTERN = re.compile(
    r"\b(?:on top of|to the (?:right|left|front|back) of|next to|between|from|on|in|to|inside)\b"
)


def _find_verb_positions(instruction: str) -> List[tuple]:
    """Find all verb occurrences and their positions in the instruction."""
    found = []
    lower = instruction.lower()
    for verb_text, verb_id in VERB_PATTERNS:
        start = 0
        while True:
            idx = lower.find(verb_text, start)
            if idx == -1:
                break
            # Ensure word boundary: not preceded by a letter
            if idx > 0 and lower[idx - 1].isalpha():
                start = idx + 1
                continue
            found.append((idx, verb_text, verb_id))
            start = idx + len(verb_text)
    # Sort by position
    found.sort(key=lambda x: x[0])
    return found


def _split_object_target(clause: str) -> tuple:
    """Split a clause into (object_str, prep_str, target_str) using preposition."""
    m = PREP_PATTERN.search(clause)
    if m:
        obj = clause[: m.start()].strip()
        prep = m.group()
        tgt = clause[m.end() :].strip()
        return obj, prep, tgt
    return clause.strip(), "", ""


def _clean_noun(text: str) -> str:
    """Clean up a noun phrase: strip articles and extra whitespace."""
    text = text.strip()
    # Remove leading "the "
    text = re.sub(r"^the\s+", "", text)
    # Remove leading "both "
    text = re.sub(r"^both\s+", "", text)
    # Collapse whitespace
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def _clean_all_articles(text: str) -> str:
    """Strip all occurrences of 'the' articles (not just leading)."""
    text = re.sub(r"\bthe\s+", "", text)
    return re.sub(r"\s+", " ", text).strip()


def parse_instruction(instruction: str) -> List[Subgoal]:
    """Parse a LIBERO instruction string into an ordered list of Subgoals."""
    lower = instruction.lower().strip()
    verb_positions = _find_verb_positions(lower)

    if not verb_positions:
        raise ValueError(f"No verbs found in instruction: {instruction}")

    # Split into clauses based on verb positions
    clauses = []
    for i, (pos, verb_text, verb_id) in enumerate(verb_positions):
        clause_start = pos + len(verb_text)
        if i + 1 < len(verb_positions):
            # Next clause starts at next verb; find the "and" connector before it
            next_pos = verb_positions[i + 1][0]
            clause_end = next_pos
            # Strip trailing " and " from clause
            raw = lower[clause_start:clause_end]
            raw = re.sub(r"\s+and\s*$", "", raw)
        else:
            raw = lower[clause_start:]
        clauses.append((verb_id, raw.strip()))

    # Process clauses into subgoals with context tracking
    subgoals: List[Subgoal] = []
    context = {"last_object": "", "last_target": ""}

    for verb_id, raw_clause in clauses:
        new_subgoals = _process_clause(verb_id, raw_clause, context)
        subgoals.extend(new_subgoals)
        # Update context from the last produced subgoal
        if new_subgoals:
            last = new_subgoals[-1]
            if last.object:
                context["last_object"] = last.object
            if last.target:
                context["last_target"] = last.target

    return subgoals


def _resolve_pronoun(text: str, context: dict, role: str) -> str:
    """Resolve pronouns ('it') and implicit references from context.

    Fallback logic: if the preferred context field is empty, try the other one.
    e.g. 'turn on the stove and put X on it' → 'it' as target should resolve
    to 'stove' (last_object) since turn_on has no target.
    """
    cleaned = _clean_noun(text)
    if cleaned in ("it", "them"):
        if role == "object":
            return context.get("last_object", "") or context.get("last_target", cleaned)
        else:
            return context.get("last_target", "") or context.get("last_object", cleaned)
    if not cleaned:
        if role == "object":
            return context.get("last_object", "") or context.get("last_target", "")
        else:
            return context.get("last_target", "") or context.get("last_object", "")
    return cleaned


def _process_clause(
    verb_id: str, raw_clause: str, context: dict
) -> List[Subgoal]:
    """Process a single verb clause into subgoal(s)."""

    # --- "put both X and Y in Z" ---
    if verb_id == "put" and "both" in raw_clause:
        return _handle_put_both(raw_clause, context)

    # --- "put X on/in Y" → pick + place ---
    if verb_id == "put":
        obj_str, prep, tgt_str = _split_object_target(raw_clause)
        obj = _clean_noun(obj_str)
        tgt = _resolve_pronoun(tgt_str, context, "target") if tgt_str else ""
        relation = prep
        # Handle "inside" without explicit target
        if not tgt and "inside" in raw_clause:
            tgt = context.get("last_target", "") or context.get("last_object", "")
            relation = "inside"
        return [
            Subgoal(type="pick", object=obj, target=""),
            Subgoal(type="place", object=obj, target=tgt, relation=relation),
        ]

    # --- "pick up X" → pick only ---
    if verb_id == "pick_up":
        obj_str, prep, tgt_str = _split_object_target(raw_clause)
        obj = _clean_noun(obj_str)
        descriptor = ""
        if prep and tgt_str:
            descriptor = f"{prep} {_clean_all_articles(tgt_str)}"
        return [Subgoal(type="pick", object=obj, descriptor=descriptor)]

    # --- "place X on/in Y" → place only ---
    if verb_id == "place":
        obj_str, prep, tgt_str = _split_object_target(raw_clause)
        obj = _resolve_pronoun(obj_str, context, "object")
        tgt = _clean_noun(tgt_str) if tgt_str else ""
        relation = prep
        return [Subgoal(type="place", object=obj, target=tgt, relation=relation)]

    # --- "open X" ---
    if verb_id == "open":
        obj = _clean_noun(raw_clause)
        return [Subgoal(type="open", object=obj)]

    # --- "close X" or "close it" ---
    if verb_id == "close":
        obj = _resolve_pronoun(raw_clause, context, "target")
        return [Subgoal(type="close", object=obj)]

    # --- "turn on X" ---
    if verb_id == "turn_on":
        obj = _clean_noun(raw_clause)
        return [Subgoal(type="turn_on", object=obj)]

    # --- "turn off X" ---
    if verb_id == "turn_off":
        obj = _clean_noun(raw_clause)
        return [Subgoal(type="turn_off", object=obj)]

    # --- "push X to Y" ---
    if verb_id == "push":
        obj_str, prep, tgt_str = _split_object_target(raw_clause)
        obj = _clean_noun(obj_str)
        tgt = _clean_noun(tgt_str) if tgt_str else ""
        relation = prep
        return [Subgoal(type="push", object=obj, target=tgt, relation=relation)]

    # --- "stack" → treat as place ---
    if verb_id == "stack":
        obj_str, prep, tgt_str = _split_object_target(raw_clause)
        obj = _clean_noun(obj_str)
        tgt = _clean_noun(tgt_str) if tgt_str else ""
        relation = prep
        return [Subgoal(type="place", object=obj, target=tgt, relation=relation)]

    raise ValueError(f"Unknown verb_id: {verb_id}")


def _handle_put_both(raw_clause: str, context: dict) -> List[Subgoal]:
    """Handle 'put both X and Y in/on Z' pattern."""
    # Remove "both"
    clause = raw_clause.replace("both", "").strip()

    # Split by preposition to get objects part and target
    obj_str, prep, tgt_str = _split_object_target(clause)
    tgt = _clean_noun(tgt_str)
    relation = prep

    # Split objects by "and"
    # Careful: need to find the "and" that separates objects, not part of a name
    parts = re.split(r"\s+and\s+", obj_str)
    if len(parts) == 2:
        obj1 = _clean_noun(parts[0])
        obj2 = _clean_noun(parts[1])
    elif len(parts) == 1:
        # "both" without "and" → plural noun referring to 2 items
        # e.g. "put both moka pots on the stove"
        name = _clean_noun(parts[0])
        # Simple singularization: strip trailing 's' (pots → pot)
        if name.endswith("s") and not name.endswith("ss"):
            name = name[:-1]
        obj1 = name
        obj2 = name
    else:
        # 3+ parts from "and" split — just take first and last
        obj1 = _clean_noun(parts[0])
        obj2 = _clean_noun(parts[-1])

    subgoals = [
        Subgoal(type="pick", object=obj1),
        Subgoal(type="place", object=obj1, target=tgt, relation=relation),
        Subgoal(type="pick", object=obj2),
        Subgoal(type="place", object=obj2, target=tgt, relation=relation),
    ]
    return subgoals


if __name__ == "__main__":
    # Quick demo with a few instructions
    test_instructions = [
        "pick up the orange juice and place it in the basket",
        "put the bowl on the plate",
        "open the top drawer and put the bowl inside",
        "put both the alphabet soup and the cream cheese box in the basket",
        "turn on the stove and put the moka pot on it",
        "put the black bowl in the bottom drawer of the cabinet and close it",
    ]
    for inst in test_instructions:
        subgoals = parse_instruction(inst)
        labels = [sg.label() for sg in subgoals]
        print(f"{inst!r}")
        print(f"  -> {labels}")
        print()