plancraft 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

plancraft/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

plancraft/evaluator.py

@@ -0,0 +1,273 @@
+import json
+import os
+import random
+import string
+import time
+
+import imageio
+import pandas as pd
+import torch
+import wandb
+from loguru import logger
+from tqdm import tqdm
+
+from plancraft.config import EvalConfig, PlancraftExample
+from plancraft.environments.actions import StopAction
+from plancraft.environments.env_real import RealPlancraft
+from plancraft.environments.env_symbolic import SymbolicPlancraft
+from plancraft.models import get_model
+
+wandb.require("core")
+
+
+class Evaluator:
+    """
+    The evaluator class handles the environment loop and model interaction
+
+    The environment is created based on the configuration and the examples are loaded from the dataset.
+    """
+
+    def __init__(self, cfg: EvalConfig):
+        self.cfg = cfg
+        self.output_dir = (
+            f"{cfg.plancraft.output_dir}/{self.evaluator_name()}/{cfg.plancraft.split}"
+        )
+        self.generation_number = 0
+
+        self.examples: list[PlancraftExample] = self.load_dataset(cfg.plancraft.split)
+
+        self.environment = self.create_env(cfg)
+        self.model = get_model(cfg)
+
+        self.record_frames = not (cfg.plancraft.environment.symbolic)
+
+        # no_op action
+        self.no_op = self.environment.action_space.no_op()
+
+    def evaluator_name(self) -> str:
+        symb_str = "real"
+        if self.cfg.plancraft.environment.symbolic:
+            symb_str = "symb"
+
+        if self.cfg.plancraft.use_maskrcnn:
+            symb_str += "_mrcnn"
+
+        model_name = self.cfg.plancraft.model.split("/")[-1]
+        if self.cfg.plancraft.adapter != "":
+            model_name = self.cfg.plancraft.adapter.split("/")[-1]
+
+        mode = self.cfg.plancraft.mode
+        if mode in ["dummy", "oracle"]:
+            return f"{mode}_{symb_str}"
+
+        actions = "|".join(self.cfg.plancraft.valid_actions)
+        return f"{self.cfg.plancraft.mode}_{symb_str}_{model_name}_{actions}"
+
+    def save_results_dict(self, example: PlancraftExample, results_dict: dict):
+        output_dir = f"{self.output_dir}/{self.generation_number}"
+        os.makedirs(output_dir, exist_ok=True)
+        json_path = f"{output_dir}/{example.id}.json"
+        with open(json_path, "w") as f:
+            json.dump(results_dict, f, indent=4)
+        wandb.save(json_path, policy="now")
+
+    def save_images(self, example: PlancraftExample, frames: list):
+        if len(frames) == 0:
+            return
+        output_dir = f"{self.output_dir}/{self.generation_number}"
+        os.makedirs(output_dir, exist_ok=True)
+        imageio.mimsave(f"{output_dir}/{example.id}.gif", frames)
+        # upload to wandb
+        wandb.save(f"{output_dir}/{example.id}.gif", policy="now")
+
+    def load_results_dict(self, example: PlancraftExample) -> dict:
+        path = f"{self.output_dir}/{self.generation_number}/{example.id}.json"
+        if not os.path.exists(path) or not self.cfg.plancraft.resume:
+            return None
+        with open(path, "r") as f:
+            return json.load(f)
+
+    def create_env(self, cfg: EvalConfig) -> RealPlancraft | SymbolicPlancraft:
+        if cfg.plancraft.environment.symbolic:
+            return SymbolicPlancraft(inventory=[])
+        return RealPlancraft(
+            inventory=[],
+            symbolic_action_space=cfg.plancraft.environment.symbolic_action_space,
+            symbolic_observation_space=cfg.plancraft.environment.symbolic_observation_space,
+            preferred_spawn_biome=cfg.plancraft.environment.preferred_spawn_biome,
+            resolution=cfg.plancraft.environment.resolution,
+        )
+
+    def close(self):
+        self.environment.close()
+
+    def load_dataset(self, dataset_split: str) -> list[PlancraftExample]:
+        with open(f"data/{dataset_split}.json", "r") as f:
+            dataset = json.load(f)
+            return [PlancraftExample(**example) for example in dataset]
+
+    def reset(
+        self,
+        example: PlancraftExample,
+    ):
+        current_inventory = example.slotted_inventory
+        self.environment.fast_reset(new_inventory=current_inventory)
+        # do a no op to an initial observation
+        obs, _, _, _ = self.environment.step(self.no_op)
+        # assert that the inventory is correct
+        if "inventory" in obs:
+            for item in current_inventory:
+                slot = item["slot"]
+                if (
+                    obs["inventory"][slot]["type"] != item["type"]
+                    or obs["inventory"][slot]["quantity"] != item["quantity"]
+                ) and item["type"] != "air":
+                    logger.warning(f"Inventory does not match expected for slot {slot}")
+                    logger.warning(f"Expected {item}")
+                    logger.warning(f"Got {obs['inventory'][slot]}")
+                    # try again
+                    self.reset(example)
+
+        objective = f"Craft an item of type: {example.target}"
+        self.model.reset_history(objective=objective)
+
+    def check_done(self, inventory: list[dict[str, int]], target: str):
+        """
+        Check that target object is obtained
+        """
+        for item in inventory:
+            if target == item["type"]:
+                # ensure item is taken out of crafting slot
+                if "slot" in item and item["slot"] != 0:
+                    return True
+                if "index" in item and item["index"] != 0:
+                    return True
+        return False
+
+    @torch.no_grad()
+    def eval_all_examples(self, progress_bar=False) -> list:
+        results = []
+        action = self.no_op.copy()
+
+        pbar = tqdm(
+            total=len(self.examples),
+            disable=not progress_bar,
+        )
+        correct = 0
+        count = 0
+
+        for example in self.examples:
+            if resume_result := self.load_results_dict(example):
+                pbar.update(self.cfg.plancraft.max_steps)
+                results.append(resume_result)
+                continue
+
+            success = False
+
+            self.reset(example)
+            action = self.no_op.copy()
+
+            while (
+                not self.model.history.check_stuck()
+                and self.model.history.num_steps < self.cfg.plancraft.max_steps
+            ):
+                # if the action is stop then we end the episode
+                if isinstance(action, StopAction):
+                    # if the action is stop and task is impossible then success
+                    # otherwise we should not have stopped
+                    success = example.impossible
+                    break
+
+                # step action
+                observation, _, _, _ = self.environment.step(action)
+
+                # check if the episode is done
+                success = self.check_done(observation["inventory"], example.target)
+                # exit if success
+                if success:
+                    break
+
+                # predict next action
+                action = self.model.step(observation)
+
+            # save results and reset
+            result = {
+                "success": success,
+                "recipe_type": example.recipe_type,
+                "number_of_steps": self.model.history.num_steps,
+                "model_trace": self.model.history.trace(),
+                "example_id": example.id,
+                "impossible": example.impossible,
+            }
+            results.append(result)
+            self.save_results_dict(example, result)
+            self.save_images(example, self.model.history.images)
+
+            correct += int(result["success"])
+            count += 1
+
+            acc = correct / count
+            pbar.set_postfix(correct=correct, count=count, acc=acc)
+            pbar.update(1)
+
+        return results
+
+    def eval_all(self):
+        logger.info(
+            f"Running evaluation over {len(self.examples)} examples {self.cfg.plancraft.num_generations} times."
+        )
+        run_name = (
+            f"{self.evaluator_name()} {self.cfg.plancraft.split}".replace(" ", "_")
+            .replace(".", "_")
+            .strip()
+        )
+
+        for n in range(self.cfg.plancraft.num_generations):
+            logger.info(f"Generation {n+1}/{self.cfg.plancraft.num_generations}")
+            run_id = "".join(random.choices(string.ascii_lowercase, k=5))
+            generation_run_name = run_name + f"_{run_id}"
+
+            wandb.init(
+                name=generation_run_name,
+                project=self.cfg.wandb.project,
+                entity=self.cfg.wandb.entity,
+                mode=self.cfg.wandb.mode,
+                group=self.cfg.plancraft.model,
+                job_type=self.cfg.plancraft.mode,
+                config=self.cfg.model_dump(),
+            )
+            time_now = time.time()
+
+            results_list = self.eval_all_examples(progress_bar=True)
+
+            results_df = pd.DataFrame(results_list)
+
+            output = {
+                "avg_success_rate": results_df["success"].mean(),
+                "avg_number_of_steps": results_df["number_of_steps"].mean(),
+                "avg_num_tokens_used": results_df["model_trace"]
+                .apply(pd.Series)["tokens_used"]
+                .mean(),
+            }
+
+            # calculate success rate for each recipe type
+            recipe_types = results_df["recipe_type"].unique()
+            for recipe_type in recipe_types:
+                mask = results_df["recipe_type"] == recipe_type
+                success_rate = results_df[mask]["success"].mean()
+                output[f"{recipe_type}_success_rate"] = success_rate
+
+            time_elapsed = time.time() - time_now
+            logger.info(f"Time elapsed: {time_elapsed:.2f}s")
+
+            logger.info(output)
+            wandb.log(output)
+            table = wandb.Table(
+                dataframe=results_df[["success", "number_of_steps", "example_id"]]
+            )
+            wandb.log({"results": table})
+            wandb.finish()
+
+            self.generation_number += 1
+
+        logger.info("Done")

plancraft/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)

plancraft/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()