plancraft 0.1.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

plancraft/evaluator.py

@@ -1,23 +1,26 @@
 import json
 import os
 import random
+import re
 import string
 import time
 
 import imageio
 import pandas as pd
-import torch
-import wandb
 from loguru import logger
 from tqdm import tqdm
 
+import wandb
 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.environment.actions import MoveAction, SmeltAction, StopAction
+from plancraft.environment.env import (
+    PlancraftEnvironment,
+    get_objective_str,
+    target_and_inventory_to_text_obs,
+)
+from plancraft.environment.search import gold_search_recipe
 from plancraft.models import get_model
-
-wandb.require("core")
+from plancraft.utils import History
 
 
 class Evaluator:
@@ -25,6 +28,11 @@ 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.
+
+    The Evaluator uses the dataset examples and initializes the environment with the example's inventory.
+
+    It is also responsible for early stopping and verifying the target object has been craft.
+    Finally, it also saves the results of the evaluation and the images generated during the evaluation.
     """
 
     def __init__(self, cfg: EvalConfig):
@@ -34,23 +42,41 @@ class Evaluator:
         )
         self.generation_number = 0
 
+        # load all examples
         self.examples: list[PlancraftExample] = self.load_dataset(cfg.plancraft.split)
 
-        self.environment = self.create_env(cfg)
-        self.model = get_model(cfg)
+        # start environment
+        self.environment = PlancraftEnvironment(
+            inventory=[],
+            resolution=cfg.plancraft.environment.resolution,
+        )
 
-        self.record_frames = not (cfg.plancraft.environment.symbolic)
+        # initialise history/dialogue tracking
+        self.history = History(
+            valid_actions=cfg.plancraft.valid_actions,
+            use_multimodal_content_format=cfg.plancraft.use_multimodal_content_format,
+            use_images=cfg.plancraft.use_images,
+            use_text_inventory=cfg.plancraft.use_text_inventory,
+            resolution=cfg.plancraft.environment.resolution,
+        )
 
-        # no_op action
-        self.no_op = self.environment.action_space.no_op()
+        # load model
+        self.model = get_model(cfg)
 
     def evaluator_name(self) -> str:
-        symb_str = "real"
-        if self.cfg.plancraft.environment.symbolic:
-            symb_str = "symb"
+        if self.cfg.plancraft.use_text_inventory and self.cfg.plancraft.use_images:
+            name_str = "both"
+        elif self.cfg.plancraft.use_images:
+            name_str = "images"
+        elif self.cfg.plancraft.use_text_inventory:
+            name_str = "text"
+        else:
+            raise ValueError(
+                "At least one of use_text_inventory or use_images should be True"
+            )
 
-        if self.cfg.plancraft.use_maskrcnn:
-            symb_str += "_mrcnn"
+        if self.cfg.plancraft.use_fasterrcnn:
+            name_str += "_fasterrcnn"
 
         model_name = self.cfg.plancraft.model.split("/")[-1]
         if self.cfg.plancraft.adapter != "":
@@ -58,10 +84,22 @@ class Evaluator:
 
         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}"
+            return f"{mode}_{name_str}"
+
+        valid_actions_to_str = {
+            "move": "m",
+            "smelt": "s",
+            "think": "t",
+            "search": "se",
+            "impossible": "i",
+        }
+        actions = "|".join(
+            [
+                valid_actions_to_str[action]
+                for action in self.cfg.plancraft.valid_actions
+            ]
+        )
+        return f"{self.cfg.plancraft.mode}_{name_str}_{model_name}_{actions}"
 
     def save_results_dict(self, example: PlancraftExample, results_dict: dict):
         output_dir = f"{self.output_dir}/{self.generation_number}"
@@ -69,7 +107,9 @@ class Evaluator:
         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")
+
+        if wandb.run is not None:
+            wandb.save(json_path, policy="now")
 
     def save_images(self, example: PlancraftExample, frames: list):
         if len(frames) == 0:
@@ -78,7 +118,8 @@ class Evaluator:
         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")
+        if wandb.run is not None:
+            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"
@@ -87,22 +128,9 @@ class Evaluator:
         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:
+        folder = os.path.dirname(os.path.abspath(__file__))
+        with open(f"{folder}/data/{dataset_split}.json", "r") as f:
             dataset = json.load(f)
             return [PlancraftExample(**example) for example in dataset]
 
@@ -110,98 +138,199 @@ class Evaluator:
         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):
+        self.environment.reset(new_inventory=example.slotted_inventory)
+        self.model.reset()
+        self.history.reset()
+
+    def check_done(self, inventory: dict, 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
+        for slot, item in inventory.items():
+            # ensure the target is in the inventory (not in slot 0)
+            if target == item["type"] and slot != 0:
+                return True
         return False
 
-    @torch.no_grad()
+    def parse_raw_model_response(
+        self, content: str
+    ) -> str | MoveAction | SmeltAction | StopAction:
+        """
+        Given a message and set of valid actions, parse the content to return the action
+        or a message if the action is not valid/requires message response
+        """
+
+        action_match = re.search(
+            f"({'|'.join(self.cfg.plancraft.valid_actions)}):", content
+        )
+        if action_match:
+            action = action_match.group(1)
+            if action == "think":
+                return "Ok"
+            elif action == "impossible":
+                reason = re.search(r"impossible: (.*)", content).group(1)
+                return StopAction(reason=reason)
+            elif action == "search":
+                search_target = re.search(r"search: (\w+)", content).group(1)
+                return gold_search_recipe(search_target)
+            else:
+                try:
+                    slot_from = re.search(r" from (\[[ABCI]?\d+\])", content).group(1)
+                    slot_to = re.search(r" to (\[[ABCI]?\d+\])", content).group(1)
+                    quantity = re.search(r"with quantity (\d+)", content).group(1)
+                    if action == "move":
+                        action = MoveAction(
+                            slot_from=slot_from,
+                            slot_to=slot_to,
+                            quantity=quantity,
+                        )
+                    else:
+                        action = SmeltAction(
+                            slot_from=slot_from,
+                            slot_to=slot_to,
+                            quantity=quantity,
+                        )
+                    return action
+                except AttributeError as e:
+                    return f"Format Error: {e}"
+        return f"Only select actions from the following: {', '.join(self.cfg.plancraft.valid_actions)}"
+
+    def convert_observation_to_message(
+        self,
+        observation: dict,
+    ) -> str | dict:
+        """
+        Convert an environment observation to the message format used by an LLM chat model
+
+        Parameters:
+        - observation: dict - The observation to convert.
+        - use_text_inventory: bool - Whether to use text inventory.
+        - use_multimodal_content_format: bool - Whether to use multimodal content format.
+        - use_images: bool - Whether to append an image to the message content - must be used with use_multimodal_content_format.
+        """
+        if self.cfg.plancraft.use_fasterrcnn:
+            # convert image to inventory using fasterrcnn
+            inventory = self.model.bbox_model.get_inventory(observation["image"].copy())
+            text_content = target_and_inventory_to_text_obs(
+                observation["target"], inventory
+            )
+        elif not self.cfg.plancraft.use_text_inventory:
+            text_content = get_objective_str(observation["target"])
+        else:
+            # if not multimodal, we only have text - we format the inventory as text
+            text_content = target_and_inventory_to_text_obs(
+                observation["target"], observation["inventory"]
+            )
+        if not self.cfg.plancraft.use_multimodal_content_format:
+            return text_content
+
+        content_list = [{"type": "text", "text": text_content}]
+        if self.cfg.plancraft.use_images:
+            content_list.append({"type": "image"})
+        return {"content": content_list}
+
+    def eval_example(self, example: PlancraftExample) -> dict:
+        """
+        Given the loaded model and an example from Plancraft
+        run the episode until success or termination.
+        Termination can happen from: early stopping (stuck) / max_steps / stop_action
+        """
+        success = False
+        num_non_env_actions = 0
+        self.reset(example)
+        action = None
+
+        # run episode until stuck or until max steps is reached
+        while (
+            not self.history.check_stuck()
+            and self.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
+            # action is external tool then it is str
+            # limit the number of consecutive non-env actions to 3
+            elif isinstance(action, str) and num_non_env_actions < 3:
+                observation = {"message": action}
+                num_non_env_actions += 1
+            # action is environment action
+            else:
+                # add action to history
+                # TODO: fix the next two lines being triggered with a str response if num_non_env_actions >= 3
+                if isinstance(action, str):
+                    observation = self.environment.step()
+                else:
+                    self.history.add_action_to_history(action)
+                    observation = self.environment.step(action)
+
+                # convert inventory observation to text message
+                observation["target"] = example.target
+                observation["message"] = self.convert_observation_to_message(
+                    observation
+                )
+                num_non_env_actions = 0
+
+                # check if the episode is done
+                success = self.check_done(observation["inventory"], example.target)
+
+            # add observation to history
+            self.history.add_observation_to_history(observation)
+            # add observation message to history
+            self.history.add_message_to_history(
+                content=observation["message"], role="user"
+            )
+
+            # exit if success
+            if success:
+                break
+
+            # predict next action
+            raw_action = self.model.step(observation, dialogue_history=self.history)
+            # add message to history
+            self.history.add_message_to_history(content=raw_action, role="assistant")
+            # parse the raw action
+            action = self.parse_raw_model_response(raw_action)
+
+        # save results and reset
+        return {
+            "success": success,
+            "recipe_type": example.recipe_type,
+            "complexity": example.complexity,
+            "number_of_steps": self.history.num_steps,
+            "model_trace": self.history.trace(),
+            "example_id": example.id,
+            "impossible": example.impossible,
+        }
+
     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:
+            logger.debug(f"Running example {example.id}")
+
             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
+            # skip impossible examples if impossible is not in valid actions
+            if (
+                example.impossible
+                and "impossible" not in self.cfg.plancraft.valid_actions
             ):
-                # 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)
+                continue
 
-                # 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,
-            }
+            result = self.eval_example(example)
             results.append(result)
             self.save_results_dict(example, result)
-            self.save_images(example, self.model.history.images)
+            self.save_images(example, self.history.images)
 
             correct += int(result["success"])
             count += 1
@@ -212,7 +341,7 @@ class Evaluator:
 
         return results
 
-    def eval_all(self):
+    def eval_all_seeds(self):
         logger.info(
             f"Running evaluation over {len(self.examples)} examples {self.cfg.plancraft.num_generations} times."
         )
@@ -222,11 +351,11 @@ class Evaluator:
             .strip()
         )
 
+        wandb.login(key=self.cfg.env_variables.wandb_api_key)
         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,
@@ -239,7 +368,6 @@ class Evaluator:
             time_now = time.time()
 
             results_list = self.eval_all_examples(progress_bar=True)
-
             results_df = pd.DataFrame(results_list)
 
             output = {
@@ -257,16 +385,38 @@ class Evaluator:
                 success_rate = results_df[mask]["success"].mean()
                 output[f"{recipe_type}_success_rate"] = success_rate
 
+            # calculate success rate for each complexity bin
+            complexity_mapping = {
+                0: "easy",
+                1: "easy",
+                2: "medium",
+                3: "hard",
+                4: "hard",
+                5: "impossible",
+            }
+            results_df["complexity_bin"] = results_df["complexity"].map(
+                complexity_mapping
+            )
+            # when impossible is True set complexity_bin to impossible
+            results_df["complexity_bin"] = results_df["complexity_bin"].fillna(
+                "impossible"
+            )
+            for complexity_bin in results_df["complexity_bin"].unique():
+                mask = results_df["complexity_bin"] == complexity_bin
+                success_rate = results_df[mask]["success"].mean()
+                output[f"{complexity_bin}_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()
+            if wandb.run is not None:
+                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
 

plancraft/generate_dataset.py

@@ -0,0 +1,72 @@
+import random
+import concurrent.futures
+from functools import partial
+
+import json
+import numpy as np
+
+from tqdm import tqdm
+
+from plancraft.environment.recipes import RECIPES
+from plancraft.environment.sampler import construct_example
+
+
+def reset_construct_example(target, num_distractors, impossible):
+    """
+    Simple retry mechanism if planner takes too long (can happen because of random sampling of initial conditions).
+    """
+    try:
+        return construct_example(
+            target=target, num_distractors=num_distractors, impossible=impossible
+        )
+    except TimeoutError:
+        print(f"TimeoutError for {target} with {num_distractors} distractors.")
+        return reset_construct_example(
+            target=target, num_distractors=num_distractors, impossible=impossible
+        )
+
+
+def process_target(recipe_target, distractors, dataset, pbar):
+    """
+    Process a single target recipe with different numbers of distractors and impossible vs possible
+    using a thread pool executor.
+    """
+    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
+        futures = []
+        for num_distractors in distractors:
+            construct_partial_false = partial(
+                reset_construct_example, recipe_target, num_distractors, False
+            )
+            construct_partial_true = partial(
+                reset_construct_example, recipe_target, num_distractors, True
+            )
+            for _ in range(10):
+                future = executor.submit(construct_partial_false)
+                futures.append((future, 1))
+            for _ in range(5):
+                future = executor.submit(construct_partial_true)
+                futures.append((future, 1))
+
+        for future, update in futures:
+            example = future.result()
+            dataset.append(example)
+            pbar.update(update)
+
+
+if __name__ == "__main__":
+    seed = 2024
+    distractors = [4, 8, 16]
+
+    random.seed(seed)
+    np.random.seed(seed)
+
+    dataset = []
+    with tqdm(total=len(RECIPES) * len(distractors) * 15) as pbar:
+        for recipe_target in list(RECIPES.keys()):
+            if len(RECIPES[recipe_target]) == 0:
+                continue
+            process_target(recipe_target, distractors, dataset, pbar)
+
+    # save the dataset
+    with open("dataset.json", "w") as file:
+        json.dump(dataset, file, indent=4)

plancraft/models/__init__.py

@@ -1,21 +1,16 @@
-from plancraft.models.base import ABCModel
-
 from plancraft.config import EvalConfig
+from plancraft.models.act import ActModel
+from plancraft.models.base import PlancraftBaseModel
 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:
+def get_model(cfg: EvalConfig) -> PlancraftBaseModel:
     """
-    Factory get model (default: ReactModel)
+    Factory get model (default: ActModel)
     """
     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)
+    return ActModel(cfg)