alfworld household skill v4

Skill library usage guide
==========================

This module exposes six task functions. Each takes a World instance `w` as
the first argument, plus one or two TYPE names (never instance names like
"apple 2" -- just "apple"). Each returns True/False for task success.

Mapping a task sentence to a call:
  - Find the two/three key nouns in the sentence. The "obj" is the thing
    being manipulated, the "dest" is the receptacle it ends up in/on, and
    for lamp tasks the "lamp" is the light source (desklamp/floorlamp).
  - Type names are lowercase, singular-ish, no spaces (e.g. "saltshaker",
    "creditcard", "coffeetable", "diningtable", "winebottle"), matching
    exactly how the game names them (drop leading articles and numbers).

1. pick_and_place(w, obj, dest)
   "put a <obj> in/on <dest>"
   Example: pick_and_place(w, "pillow", "sofa")
     -- task: "put a pillow on the sofa"

2. clean_and_place(w, obj, dest)
   "put a clean <obj> in <dest>" / "clean some <obj> and put it in <dest>"
   Example: clean_and_place(w, "lettuce", "diningtable")
     -- task: "put a clean lettuce in diningtable"

3. heat_and_place(w, obj, dest)
   "heat some <obj> and put it in <dest>" / "put a hot <obj> in <dest>"
   Example: heat_and_place(w, "egg", "countertop")
     -- task: "heat some egg and put it in countertop"

4. cool_and_place(w, obj, dest)
   "cool some <obj> and put it in <dest>" / "put a cold <obj> in <dest>"
   Example: cool_and_place(w, "potato", "microwave")
     -- task: "cool some potato and put it in microwave"

5. examine_in_light(w, obj, lamp)
   "examine the <obj> with the <lamp>" / "look at <obj> under the <lamp>"
   Example: examine_in_light(w, "book", "desklamp")
     -- task: "examine the book with the desklamp"
   (lamp is usually "desklamp" or "floorlamp")

6. pick_two_and_place(w, obj, dest)
   "put two <obj> in/on <dest>" / "find two <obj> and put them in <dest>"
   Example: pick_two_and_place(w, "cd", "safe")
     -- task: "put two cds in safe"
   Example: pick_two_and_place(w, "soapbottle", "toilet")
     -- task: "find two soapbottle and put them in toilet"

Notes:
  - obj/dest/lamp must be singular type names as they appear in the game's
    object list (e.g. "soapbar" not "soap bar", "cellphone" not "cell phone").
    Strip leading articles ("a", "an", "the") and any trailing numbers.
  - If a sentence mentions two actions (e.g. "clean" + "put"), pick the
    matching combined function (clean_and_place / heat_and_place /
    cool_and_place) rather than calling pick_and_place twice.
  - Every function is safe to call even if the agent is already holding an
    unrelated object (it will be set down first), or if the target object
    was already moved/placed by an earlier partial attempt (search checks
    the current spot, previously visited receptacles, and the destination
    itself before falling back to the full room search order).
  - pick_two_and_place in particular never gets stuck holding an object: if
    a placement attempt fails it puts the item down again before continuing
    to search for the next instance, since only one object can be held at a
    time.
  - Every function already handles opening closed receptacles/appliances
    and stays within the environment's ~40 step budget.

"""Task-level skills built on top of primitives.World.

Each function below performs one ALFWorld task type end-to-end using only
World's verbs (go/open/close/take/put/clean/heat/cool/use/examine) and
queries (objects_here/receptacles_of/instances_of/receptacles), plus the
externally supplied `search_order(obj_type, receptacles)` prior that orders
receptacle instances by how likely they are to hold a given object type.

No task logic (search order, plans) lives in primitives.py; it all lives here.

Robustness notes:
  - Every entry point may be called while the agent is already holding a
    (possibly wrong) object, and/or after earlier partial progress (target
    already at destination, receptacles already opened/visited). Each
    function starts by getting rid of a wrong held item and by searching
    broadly enough (current location, previously visited receptacles, the
    destination, then the full prior order) to find the target wherever it
    ended up.
  - pick_two_and_place additionally must never get "stuck" holding an item:
    if a placement attempt fails, the held item is set down again before
    continuing the search, since the game only allows holding one object at
    a time and a stuck-holding bug would silently stall all further takes.
"""

from __future__ import annotations

# Soft step budget: we try to leave a few steps of slack under the ~40 step
# limit so a task never gets cut off mid-action.
MAX_STEPS = 40
SAFETY_MARGIN = 4


def _steps_left(w) -> int:
    try:
        return MAX_STEPS - w.steps
    except Exception:
        return MAX_STEPS


def _budget_ok(w, cost: int = 2) -> bool:
    return _steps_left(w) > SAFETY_MARGIN + cost


# ---------------------------------------------------------------------------
# type / instance helpers
# ---------------------------------------------------------------------------

def _norm_type(t: str) -> str:
    """Normalize a caller-supplied type name for comparison purposes."""
    return t.strip().lower().replace(" ", "")


def _inst_type(instance_name: str) -> str:
    """'soap bar 1' -> 'soapbar', 'apple' -> 'apple'."""
    if not instance_name:
        return ""
    parts = instance_name.strip().split()
    if parts and parts[-1].isdigit():
        parts = parts[:-1]
    return "".join(parts).lower()


# ---------------------------------------------------------------------------
# internal helpers
# ---------------------------------------------------------------------------

def _refresh(w, recep):
    """Go to a receptacle, opening it if it's closed. Returns the obs text."""
    obs = w.go(recep)
    if w.is_closed_here(obs):
        obs = w.open(recep)
    return obs


def _put_down_somewhere(w):
    """Drop whatever is currently held onto the current or any receptacle."""
    if not w.holding:
        return
    if w.at:
        w.put(w.holding, w.at)
        if not w.holding:
            return
    for r in w.receptacles:
        if not _budget_ok(w):
            break
        obs = w.go(r)
        if w.is_closed_here(obs):
            w.open(r)
        w.put(w.holding, r)
        if not w.holding:
            return


def _ensure_not_holding_wrong(w, obj_type):
    """If holding something that isn't obj_type, put it down first."""
    if not w.holding:
        return
    if _inst_type(w.holding) == _norm_type(obj_type):
        return
    _put_down_somewhere(w)


def _locate_next(w, obj_type, visited, taken=()):
    """Find one instance of obj_type not already in `taken`.

    Order of search:
      1. wherever we currently stand (cheap, no extra step),
      2. receptacles we've already visited this task (from cache, in case
         they held more than one instance of obj_type -- important for
         pick_two_and_place when both copies sit on the same receptacle),
      3. the remaining receptacles in search_order(obj_type, ...), each
         visited (go + open-if-closed) at most once.

    Mutates `visited` in place. Returns (instance_name, receptacle_name) or
    (None, None). Backs off early if the step budget is getting tight.
    """
    # 1. current location, no extra step needed
    if w.at:
        items = w.objects_here()
        insts = [i for i in w.instances_of(obj_type, items) if i not in taken]
        if insts:
            visited.add(w.at)
            return insts[0], w.at

    # 2. receptacles already visited (cached), re-check for leftovers
    for recep in list(visited):
        cached = w.seen.get(recep, [])
        insts = [i for i in w.instances_of(obj_type, cached) if i not in taken]
        if insts and _budget_ok(w):
            _refresh(w, recep)
            items = w.objects_here()
            insts2 = [i for i in w.instances_of(obj_type, items) if i not in taken]
            if insts2:
                return insts2[0], recep

    # 3. fresh receptacles via the prior
    for recep in search_order(obj_type, w.receptacles):
        if recep in visited:
            continue
        if not _budget_ok(w):
            break
        visited.add(recep)
        _refresh(w, recep)
        items = w.objects_here()
        insts = [i for i in w.instances_of(obj_type, items) if i not in taken]
        if insts:
            return insts[0], recep
    return None, None


def _take(w, inst, recep):
    """Take inst from recep; returns True if now holding it."""
    w.take(inst, recep)
    return w.holding == inst


def _place(w, inst, dest_type):
    """Try to put the held instance onto/into some instance of dest_type."""
    for d in w.receptacles_of(dest_type):
        if not _budget_ok(w):
            break
        obs = w.go(d)
        if w.is_closed_here(obs):
            obs = w.open(d)
        obs = w.put(inst, d)
        if obs and ("you put" in obs.lower() or "you move" in obs.lower()):
            return True
    return False


def _find_and_take(w, obj_type, visited, taken=()):
    """Locate an instance of obj_type and pick it up, retrying on failure.

    If already holding a matching instance (not in taken), use it directly
    without spending any extra steps -- this covers resumed/retried calls.
    """
    _ensure_not_holding_wrong(w, obj_type)
    if w.holding and _inst_type(w.holding) == _norm_type(obj_type) and w.holding not in taken:
        return w.holding

    failed = set()
    for _ in range(5):
        if not _budget_ok(w):
            return None
        inst, recep = _locate_next(w, obj_type, visited, tuple(taken) + tuple(failed))
        if not inst:
            return None
        if _take(w, inst, recep):
            return inst
        failed.add(inst)
    return None


def _go_appliance(w, appliance_type):
    """Go to (and open if needed) the first instance of an appliance type."""
    insts = w.receptacles_of(appliance_type)
    if not insts:
        return None
    inst = insts[0]
    obs = w.go(inst)
    if w.is_closed_here(obs):
        w.open(inst)
    return inst


# ---------------------------------------------------------------------------
# entry points
# ---------------------------------------------------------------------------

def pick_and_place(w, obj, dest):
    visited = set()
    inst = _find_and_take(w, obj, visited)
    if not inst:
        return False
    return _place(w, inst, dest)


def clean_and_place(w, obj, dest):
    visited = set()
    inst = _find_and_take(w, obj, visited)
    if not inst:
        return False

    sinks = w.receptacles_of("sinkbasin")
    if not sinks:
        return False
    sink = sinks[0]
    obs = w.go(sink)
    if w.is_closed_here(obs):
        w.open(sink)
    w.clean(inst, sink)

    return _place(w, inst, dest)


def heat_and_place(w, obj, dest):
    visited = set()
    inst = _find_and_take(w, obj, visited)
    if not inst:
        return False

    m = _go_appliance(w, "microwave")
    if not m:
        return False
    w.heat(inst, m)

    return _place(w, inst, dest)


def cool_and_place(w, obj, dest):
    visited = set()
    inst = _find_and_take(w, obj, visited)
    if not inst:
        return False

    f = _go_appliance(w, "fridge")
    if not f:
        return False
    w.cool(inst, f)

    return _place(w, inst, dest)


def examine_in_light(w, obj, lamp):
    visited = set()
    inst = _find_and_take(w, obj, visited)
    if not inst:
        return False

    linst, lrecep = _locate_next(w, lamp, visited)
    if not linst:
        return False

    if w.at != lrecep:
        w.go(lrecep)
        items = w.objects_here()
        fresh = w.instances_of(lamp, items)
        if fresh:
            linst = fresh[0]

    w.use(linst)
    return w.won


def pick_two_and_place(w, obj, dest):
    visited = set()
    taken: list[str] = []
    failed: set[str] = set()

    # Drop any wrong-type held object first so it never blocks picking up
    # instances of obj (the game only lets you hold one item at a time).
    _ensure_not_holding_wrong(w, obj)

    # Check whether one (or both) instances already sit at dest from an
    # earlier/partial attempt -- count them without moving them.
    for d in w.receptacles_of(dest):
        if not _budget_ok(w):
            break
        _refresh(w, d)
        visited.add(d)
        items = w.objects_here()
        for inst in w.instances_of(obj, items):
            if inst not in taken:
                taken.append(inst)
        if len(taken) >= 2:
            break

    # If we're currently holding a matching instance not yet counted (e.g.
    # resumed after a prior partial attempt), try to place it right away.
    if (
        len(taken) < 2
        and w.holding
        and _inst_type(w.holding) == _norm_type(obj)
        and w.holding not in taken
    ):
        held = w.holding
        if _place(w, held, dest):
            taken.append(held)
        else:
            failed.add(held)
            _put_down_somewhere(w)

    attempts = 0
    while len(taken) < 2 and attempts < 6 and _budget_ok(w, 4):
        attempts += 1
        inst, recep = _locate_next(w, obj, visited, tuple(taken) + tuple(failed))
        if not inst:
            break
        if not _take(w, inst, recep):
            failed.add(inst)
            continue
        if _place(w, inst, dest):
            taken.append(inst)
        else:
            # Placement failed -- do NOT stay stuck holding it, or every
            # future take() call will silently fail for the rest of the
            # task. Put it down somewhere and keep searching.
            failed.add(inst)
            _put_down_somewhere(w)

    return len(taken) >= 2