plancraft 0.1.0__py3-none-any.whl

environments/sampler.py

@@ -0,0 +1,224 @@
+import math
+import random
+from collections import Counter
+
+import numpy as np
+from plancraft.environments.items import all_data, ALL_ITEMS
+from plancraft.environments.recipes import RECIPES
+from plancraft.environments.planner import optimal_planner, get_ancestors
+
+
+MAX_STACK_SIZE = {}
+for data_item in all_data["items"]:
+    if data_item["stackable"]:
+        MAX_STACK_SIZE[data_item["type"]] = data_item["stackSize"]
+    else:
+        MAX_STACK_SIZE[data_item["type"]] = 1
+
+
+def sample_distractors(
+    exclude_set: set = None, num_distractors: int = 16
+) -> dict[str, int]:
+    distractors = {}
+    while len(distractors) < num_distractors:
+        item = random.choice(ALL_ITEMS)
+        if exclude_set is not None and item in exclude_set:
+            continue
+        count = random.randint(1, MAX_STACK_SIZE[item])
+        distractors[item] = count
+    return distractors
+
+
+def assign_to_slots(inventory: dict[str, int]) -> list[dict]:
+    # slots available outside of crafting interface
+    available_slots = list(range(10, 46))
+    random.shuffle(available_slots)
+    inventory_list = []
+
+    for item, total_count in inventory.items():
+        while total_count > 0:
+            if len(available_slots) == 0:
+                print("Not enough slots available")
+                break
+            slot = available_slots.pop()
+            count_in_slot = min(total_count, MAX_STACK_SIZE[item])
+            inventory_list.append({"slot": slot, "item": item, "count": count_in_slot})
+            total_count -= count_in_slot
+
+    return inventory_list
+
+
+def sample_recipes(
+    target: str,
+    overall_exclude_set: set,
+    target_count: int = 1,
+    current_depth=0,
+    max_depth=20,
+) -> tuple[set, set]:
+    # stop if the depth is too high
+    if current_depth > max_depth:
+        return {}, overall_exclude_set
+
+    # get all the recipes that can craft the target
+    overall_exclude_set.update([target])
+    local_exclude_set = set()
+    random_recipes = []
+    for r in RECIPES[target]:
+        recipe_inputs, exclude_set = r.sample_inputs()
+        # if inputs are already in the exclude set, skip this recipe (ensures no cycle)
+        if exclude_set.intersection(overall_exclude_set):
+            return {}, overall_exclude_set
+        local_exclude_set.update(exclude_set)
+        random_recipes.append((r, recipe_inputs))
+
+    overall_exclude_set |= local_exclude_set
+
+    # no recipes found
+    if len(random_recipes) == 0:
+        return {}, overall_exclude_set
+
+    # sample a random recipe
+    random_recipe = random.choice(random_recipes)
+    recipe, start_inputs = random_recipe
+
+    # recipe will not produce enough
+    if recipe.result.count < target_count:
+        # must do recipe X times
+        recipe_multiplier = math.ceil(target_count / recipe.result.count)
+        start_inputs = {k: v * recipe_multiplier for k, v in start_inputs.items()}
+
+    for input_item in list(start_inputs.keys()):
+        # randomize depth first search to end early
+        if random.choice([True, False]):
+            continue
+
+        children_recipe_inputs, updated_exclude_set = sample_recipes(
+            target=input_item,
+            overall_exclude_set=overall_exclude_set,
+            target_count=start_inputs[input_item],
+            current_depth=current_depth + 1,
+        )
+        if len(children_recipe_inputs) == 0:
+            continue
+
+        overall_exclude_set.update(updated_exclude_set)
+
+        # remove recipe input item since we are crafting it
+        start_inputs[input_item] = 0
+
+        # add the children recipe inputs
+        for item, count in children_recipe_inputs.items():
+            start_inputs[item] = start_inputs.get(item, 0) + count
+
+    overall_exclude_set = overall_exclude_set - {None}
+    start_inputs = {k: v for k, v in start_inputs.items() if v > 0}
+
+    return start_inputs, overall_exclude_set
+
+
+def remove_ancestor_items(target: str, inventory: dict[str, int]) -> dict[str, int]:
+    ancestors = set(get_ancestors(target))
+    possible_items = set(inventory.keys())
+    items_to_remove = list(ancestors.intersection(possible_items))
+    num_items = random.randint(1, len(items_to_remove))
+    for item in random.sample(items_to_remove, num_items):
+        count_to_remove = random.randint(1, inventory[item])
+        inventory[item] -= count_to_remove
+        if inventory[item] == 0:
+            del inventory[item]
+    return inventory
+
+
+def construct_example(
+    target: str,
+    num_distractors: 16,
+    impossible=False,
+) -> list[dict]:
+    """
+    For a given target object, number of distractors, and impossible flag
+    Return a dictionary with the start inventory for the crafting task
+
+    The crafting task should be to craft the target, the inventory should contain
+    the resources required for the recipe to be crafted.
+
+    The number of distractors are how many random items should be added to the inventory.
+
+    If impossible is True, the target item should not be craftable with the given inventory.
+    """
+
+    # sample the recipe
+    inventory, overall_exclude_set = sample_recipes(target, set())
+    if impossible:
+        # if impossible then remove one or more items from the inventory
+        inventory = remove_ancestor_items(
+            target,
+            inventory,
+        )
+
+    # add distractors to the inventory
+    distractors = sample_distractors(overall_exclude_set, num_distractors)
+    inventory.update(distractors)
+
+    optimal_path = optimal_planner(target, inventory)
+    # @TODO this is a hack to ensure that we don't have impossible examples
+    while optimal_path is not None and impossible:
+        inventory = remove_ancestor_items(target, inventory)
+        optimal_path = optimal_planner(target, inventory)
+
+    # assign to slots
+    inventory_list = assign_to_slots(inventory)
+    example = {
+        "inventory": inventory,
+        "slotted_inventory": inventory_list,
+        "target": target,
+        "num_distractors": num_distractors,
+        "impossible": impossible,
+    }
+    # either impossible and no path or not impossible and path exists
+    assert (impossible and optimal_path is None) or (
+        not impossible and optimal_path is not None
+    )
+
+    if not impossible:
+        example["optimal_path_length"] = len(optimal_path)
+        example["optimal_path"] = [r.result.item for (r, i) in optimal_path]
+        example["inventory_trace"] = [i for (r, i) in optimal_path]
+        items_used, unique_items_used = calculate_stats_from_inventory_trace(
+            [example["inventory"]] + example["inventory_trace"]
+        )
+        example["items_used"] = items_used
+        example["unique_items_used"] = unique_items_used
+
+    return example
+
+
+def calculate_stats_from_inventory_trace(
+    inventory_trace: list[dict],
+) -> tuple[int, int]:
+    total_items_used = 0
+    total_unique_items_used = 0
+
+    for a, b in zip(inventory_trace[:-1], inventory_trace[1:]):
+        diff = Counter(a) - Counter(b)
+        total_items_used += sum(diff.values())
+        total_unique_items_used += len(diff)
+
+    return total_items_used, total_unique_items_used
+
+
+def generate_dataset(seed=2024, distractors=[4, 8, 16], num_examples=10):
+    random.seed(seed)
+    np.random.seed(seed)
+
+    dataset = []
+    for recipe_target in list(RECIPES.keys()):
+        if len(RECIPES[recipe_target]) == 0:
+            continue
+        for num_distractors in distractors:
+            for _ in range(num_examples):
+                example = construct_example(
+                    target=recipe_target, num_distractors=num_distractors
+                )
+                dataset.append(example)
+
+    return dataset

models/__init__.py

@@ -0,0 +1,21 @@
+from plancraft.models.base import ABCModel
+
+from plancraft.config import EvalConfig
+from plancraft.models.dummy import DummyModel
+from plancraft.models.react import ReactModel
+from plancraft.models.oracle import OracleModel
+from plancraft.models.act import ActModel
+
+
+def get_model(cfg: EvalConfig) -> ABCModel:
+    """
+    Factory get model (default: ReactModel)
+    """
+    if cfg.plancraft.mode == "dummy":
+        return DummyModel(cfg)
+    elif cfg.plancraft.mode == "oracle":
+        return OracleModel(cfg)
+    elif cfg.plancraft.mode == "act":
+        return ActModel(cfg)
+    else:
+        return ReactModel(cfg)

models/act.py

@@ -0,0 +1,184 @@
+import copy
+import torch
+from dotenv import load_dotenv
+
+from plancraft.config import EvalConfig
+from plancraft.environments.actions import (
+    NoOp,
+    StopAction,
+    SymbolicAction,
+)
+from plancraft.models.base import ABCModel, History
+from plancraft.models.bbox_model import IntegratedBoundingBoxModel
+from plancraft.models.few_shot_images import load_prompt_images
+from plancraft.models.generators import (
+    OAMGenerator,
+    OpenAIGenerator,
+    TransformersGenerator,
+)
+from plancraft.models.prompts import get_prompt_example, get_system_prompt
+from plancraft.models.utils import (
+    convert_observation_to_message,
+    parse_content_response,
+)
+
+
+load_dotenv()
+
+
+class ActModel(ABCModel):
+    """
+    Model that does action without thinking step
+    """
+
+    def __init__(self, cfg: EvalConfig):
+        assert (
+            cfg.plancraft.environment.symbolic_action_space
+        ), "Real action space unsupported"
+        self.cfg = cfg
+        self.env_is_multimodal = not cfg.plancraft.environment.symbolic
+        self.use_maskrcnn = cfg.plancraft.use_maskrcnn
+        self.use_multimodal_content_format = cfg.plancraft.use_multimodal_content_format
+        self.use_text_inventory = cfg.plancraft.use_text_inventory
+        self.use_images = cfg.plancraft.use_images
+
+        self.bbox_model = None
+        if self.use_maskrcnn:
+            assert self.env_is_multimodal, "MaskRCNN only supported in multimodal mode"
+            self.bbox_model = IntegratedBoundingBoxModel.from_pretrained(
+                "gautierdag/plancraft-maskrcnn"
+            )
+            self.bbox_model.eval()
+            if torch.cuda.is_available():
+                self.bbox_model.cuda()
+            # MaskRCNN is not multimodal model but a separate model
+
+        self.few_shot = cfg.plancraft.few_shot
+        self.use_system_prompt = cfg.plancraft.system_prompt
+        self.max_invalid_actions = 3
+
+        # underlying language model
+        if "gpt-4o" in cfg.plancraft.model:
+            self.use_multimodal_content_format = True
+            self.llm = OpenAIGenerator(
+                use_images=self.use_images, model_name=cfg.plancraft.model
+            )
+        elif "oam" in cfg.plancraft.model:
+            self.llm = OAMGenerator(model_name=cfg.plancraft.model)
+        else:
+            # model is transformers based
+            self.llm = TransformersGenerator(
+                model_name=cfg.plancraft.model,
+                tokenizer_name=cfg.plancraft.tokenizer,
+                quantize=cfg.plancraft.quantize,
+                use_hot_cache=cfg.plancraft.hot_cache,
+                adapter_name=cfg.plancraft.adapter,
+            )
+
+        self.prompt_images = []
+
+        self.valid_actions = cfg.plancraft.valid_actions
+        self.system_prompt_text = get_system_prompt(self.valid_actions)
+
+        examples = get_prompt_example(
+            self.valid_actions,
+            use_text_inventory=self.use_text_inventory,
+            use_multimodal_content_format=self.use_multimodal_content_format,
+            use_images=self.use_images,
+        )
+        if self.env_is_multimodal and self.use_images:
+            self.prompt_images = load_prompt_images()
+
+        if self.use_multimodal_content_format:
+            self.system_prompt = {
+                "role": "system",
+                "content": [
+                    {"text": copy.deepcopy(self.system_prompt_text), "type": "text"}
+                ],
+            }
+        else:
+            self.system_prompt = {
+                "role": "system",
+                "content": copy.deepcopy(self.system_prompt_text),
+            }
+
+        if not self.few_shot:
+            examples = []
+        if not self.use_system_prompt:
+            self.system_prompt = None
+
+        self.history = History(
+            initial_dialogue=examples,
+            use_multimodal_content_format=self.use_multimodal_content_format,
+        )
+
+        self.max_messages_window = cfg.plancraft.max_message_window
+        self.kv_cache = None
+
+    def reset_history(
+        self,
+        objective: str,
+    ):
+        examples = []
+        if self.few_shot:
+            examples = get_prompt_example(
+                self.valid_actions,
+                use_text_inventory=self.use_text_inventory,
+                use_multimodal_content_format=self.use_multimodal_content_format,
+                use_images=self.use_images,
+            )
+
+        self.history.reset(objective=objective, initial_dialogue=examples)
+        self.llm.reset()
+
+    def step(self, observation: dict) -> SymbolicAction | StopAction:
+        self.history.add_observation_to_history(observation)
+
+        # add observation to history
+        observation_message = convert_observation_to_message(
+            observation,
+            objective=self.history.objective,
+            bbox_model=self.bbox_model,
+            oam_model="oam" in self.llm.model_name,
+            use_text_inventory=self.use_text_inventory,
+            use_multimodal_content_format=self.use_multimodal_content_format,
+            use_images=self.use_images,
+        )
+        self.history.add_message_to_history(content=observation_message, role="user")
+
+        # Iterate until valid action
+        i = 0
+        while i < self.max_invalid_actions:
+            # add observation to history
+            message_window, image_window = self.llm.prepare_messages(
+                history=self.history,
+                max_messages_window=self.max_messages_window,
+                system_prompt=self.system_prompt,
+                prompt_images=self.prompt_images,
+            )
+            action_messages, action_token_used = self.llm.generate_unconstrained(
+                batch_messages=[message_window],
+                images=[image_window],
+            )
+            self.history.tokens_used += action_token_used
+
+            action_message = action_messages[0].split("\n")[0].strip()
+
+            self.history.add_message_to_history(
+                content=action_message, role="assistant"
+            )
+            response = parse_content_response(
+                action_message, valid_actions=self.valid_actions
+            )
+            if not isinstance(response, str):
+                # valid action
+                self.history.add_action_to_history(response)
+                return response
+
+            self.history.add_message_to_history(
+                content=response,
+            )
+            i += 1
+
+        # if no action is found after max_invalid_actions, default to useless move action
+        return NoOp()

models/base.py

@@ -0,0 +1,152 @@
+import abc
+
+from copy import copy
+from collections import Counter
+
+from plancraft.environments.actions import (
+    SymbolicMoveAction,
+    RealActionInteraction,
+    SymbolicSmeltAction,
+)
+
+
+class History:
+    def __init__(
+        self,
+        objective: str = "",
+        initial_dialogue: list[dict] = [],
+        use_multimodal_content_format=False,
+    ):
+        self.dialogue_history = initial_dialogue
+        self.initial_dialogue_length = len(initial_dialogue)
+        self.action_history = []
+        self.inventory_history = []
+        self.inventory_counters = []
+        self.images = []
+        self.objective = objective
+        self.tokens_used = 0
+        self.use_multimodal_content_format = use_multimodal_content_format
+
+    def add_message_to_history(self, content: str | dict, role="user"):
+        if role == "assistant":
+            print(content)
+
+        if isinstance(content, dict):
+            assert "content" in content, "content key not found in message"
+            content["role"] = role
+            self.dialogue_history.append(content)
+        else:
+            # fix for listed content type
+            if self.use_multimodal_content_format:
+                return self.add_message_to_history(
+                    content={
+                        "content": [{"type": "text", "text": content}],
+                        "role": role,
+                    },
+                    role=role,
+                )
+            else:
+                self.dialogue_history.append({"role": role, "content": content})
+
+    def add_action_to_history(
+        self, action: SymbolicSmeltAction | RealActionInteraction | SymbolicMoveAction
+    ):
+        if action is None:
+            return
+        self.action_history.append(action.model_dump())
+
+    def add_inventory_to_history(self, inventory: list[dict[str, int]]):
+        self.inventory_history.append(inventory)
+
+        # count inventory
+        counter = Counter()
+        for item in inventory:
+            # ignore slot 0
+            if "slot" in item and item["slot"] == 0:
+                continue
+            if "index" in item and item["index"] == 0:
+                continue
+            counter[item["type"]] += item["quantity"]
+
+        self.inventory_counters.append(counter)
+
+    def add_image_to_history(self, image):
+        self.images.append(image)
+
+    def add_observation_to_history(self, observation: dict):
+        if observation is None:
+            return
+        if "inventory" in observation:
+            clean_inv = []
+            # remove empty slots
+            for item in observation["inventory"]:
+                if item["quantity"] > 0:
+                    clean_inv.append(item)
+            self.add_inventory_to_history(clean_inv)
+        if "pov" in observation:
+            self.add_image_to_history(observation["pov"])
+
+    def __str__(self):
+        return str(self.dialogue_history)
+
+    def reset(self, objective: str = "", initial_dialogue: list[dict] = []):
+        self.dialogue_history = initial_dialogue
+        self.action_history = []
+        self.inventory_history = []
+        self.inventory_counters = []
+        self.images = []
+        self.objective = objective
+
+    def set_objective(self, objective: str):
+        self.objective = objective
+
+    def trace(self):
+        return {
+            "dialogue_history": copy(
+                self.dialogue_history[self.initial_dialogue_length :]
+            ),
+            "action_history": copy(self.action_history),
+            "inventory_history": copy(self.inventory_history),
+            "objective": copy(self.objective),
+            "tokens_used": copy(self.tokens_used),
+        }
+
+    @property
+    def num_steps(self):
+        return len(self.action_history)
+
+    def check_stuck(self, max_steps_no_change: int = 10) -> bool:
+        """
+        If inventory content does not change for max_steps_no_change steps
+        the agent is considered stuck.
+
+        With N=10, the oracle solver can still solve 100% of the examples
+        """
+        if len(self.inventory_counters) <= max_steps_no_change:
+            return False
+
+        return all(
+            c == self.inventory_counters[-max_steps_no_change - 1]
+            for c in self.inventory_counters[-max_steps_no_change - 1 :]
+        )
+
+
+class ABCModel(abc.ABC):
+    """
+    Model class must implement the following methods to work with evaluator
+    """
+
+    @abc.abstractmethod
+    def step(
+        self, observation: list[dict]
+    ) -> list[SymbolicMoveAction | RealActionInteraction | SymbolicSmeltAction]:
+        """
+        Model should output a valid action based on the 3 types available
+
+        Note this is a batch operation, so the model should return a list of actions
+        for each observation in the batch
+        """
+        raise NotImplementedError()
+
+    def reset_history(self, objective: str = ""):
+        self.history.reset(objective=objective)