plancraft 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

plancraft/models/utils.py

@@ -2,22 +2,12 @@ import base64
 import glob
 import io
 import pathlib
-import re
 
 import numpy as np
 import torch
 from PIL import Image
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from plancraft.environments.actions import (
-    StopAction,
-    SymbolicAction,
-    SymbolicMoveAction,
-    SymbolicSmeltAction,
-    convert_from_slot_index,
-)
-from plancraft.environments.recipes import RECIPES
-
 
 def numpy_to_base64(img_array: np.ndarray, image_format: str = "PNG") -> str:
     """
@@ -150,140 +140,3 @@ def tokenize(
             padding=True,
         )
     return tokenized_messages
-
-
-def objective_and_inventory_to_str(objective: str, inventory: list[dict]) -> str:
-    inventory_str = ""
-    for item in inventory:
-        if item["quantity"] > 0:
-            if "index" in item:
-                slot = item["index"]
-            else:
-                slot = item["slot"]
-
-            if isinstance(slot, int):
-                slot = convert_from_slot_index(slot)
-
-        inventory_str += f"\n - {item['type']} {slot} quantity {item['quantity']}"
-
-    return f"{objective}\ninventory:{inventory_str}"
-
-
-def convert_observation_to_message(
-    observation: dict,
-    objective: str,
-    bbox_model=None,
-    oam_model=False,
-    use_text_inventory=True,
-    use_multimodal_content_format=False,
-    use_images=False,
-) -> str | dict:
-    """
-    Convert an observation to a message format
-
-    Parameters:
-    - observation: dict - The observation to convert.
-    - objective: str - The objective of the observation.
-    - bbox_model: Optional - The bounding box model to use.
-    - oam_model: bool - Whether to use the OAM model.
-    - 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 bbox_model is not None:
-        # convert to tensor
-        inventory = bbox_model.get_inventory(observation["pov"].copy())
-        text_content = objective_and_inventory_to_str(
-            objective, sorted(inventory, key=lambda x: x["slot"])
-        )
-    elif oam_model:
-        text_content = f"{objective}\ninventory:\n"
-    elif not use_text_inventory:
-        text_content = objective
-    else:
-        # if not multimodal, we only have text - we just dump a JSON of the inventory
-        inventory = []
-        for o in observation["inventory"]:
-            if o["quantity"] > 0:
-                inventory.append(
-                    {
-                        "type": o["type"],
-                        "slot": convert_from_slot_index(o["index"]),
-                        "quantity": o["quantity"],
-                    }
-                )
-        text_content = objective_and_inventory_to_str(objective, inventory)
-
-    if not use_multimodal_content_format:
-        return text_content
-
-    content_list = [{"type": "text", "text": text_content}]
-    if use_images:
-        content_list.append({"type": "image"})
-    return {"content": content_list}
-
-
-def gold_search_recipe(recipe_name: str) -> str:
-    """
-    Gold search recipe for the given observation and action
-    """
-    if recipe_name not in RECIPES:
-        return "Could not find a recipe by that name."
-
-    out_string = f"Recipes to craft {recipe_name}:\n"
-    for i, r in enumerate(RECIPES[recipe_name]):
-        if r.recipe_type != "smelting":
-            # sample a valid input grid (note that this is not guaranteed to be the only valid grid)
-            input_crafting_grid = r.sample_input_crafting_grid()
-            recipe_instructions = ""
-            for item in input_crafting_grid:
-                recipe_instructions += (
-                    f"{item['type']} at {convert_from_slot_index(item['slot'])}\n"
-                )
-        else:
-            # smelting recipe
-            recipe_instructions = f"smelt {r.ingredient}\n"
-        out_string += f"recipe {i+1}:\n{recipe_instructions}"
-    return out_string
-
-
-def parse_content_response(
-    content: str, valid_actions: list[str] = ["smelt", "move"]
-) -> str | SymbolicAction | 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(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 = SymbolicMoveAction(
-                        slot_from=slot_from,
-                        slot_to=slot_to,
-                        quantity=quantity,
-                    )
-                else:
-                    action = SymbolicSmeltAction(
-                        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(valid_actions)}"

plancraft/train/dataset.py

@@ -21,7 +21,7 @@ TEMPLATES = {
 class PlancraftDialogueDataset(Dataset):
     def __init__(
         self,
-        dataset_dir: str = "data/oracle",
+        dataset_dir: str = "oracle_trajectories",
         use_images=False,
         trace_mode="oa",
         split="train",

plancraft/utils.py

@@ -1,12 +1,184 @@
 import glob
 import pathlib
+from collections import Counter
+from copy import copy
 
-import cv2
-import numpy as np
 import torch
-import torch.nn.functional as F
 from loguru import logger
-from PIL import Image
+
+from plancraft.environment.actions import (
+    ActionHandlerBase,
+    MoveAction,
+    SmeltAction,
+)
+from plancraft.environment.prompts import (
+    get_system_prompt,
+    get_prompt_example,
+    load_prompt_images,
+)
+
+
+class History:
+    """
+    History class to keep track of dialogue, actions, inventory and images
+    Args:
+        valid_actions: list of valid actions names
+        initial_dialogue: list of dialogue messages
+        use_multimodal_content_format: whether to use multimodal content format (list of content with types)
+    """
+
+    def __init__(
+        self,
+        actions: list[ActionHandlerBase] = [],
+        use_multimodal_content_format=False,
+        few_shot=False,
+        use_images=False,
+        use_text_inventory=False,
+        resolution="high",
+    ):
+        self.action_handlers = actions
+        self.use_multimodal_content_format = use_multimodal_content_format
+        self.few_shot = few_shot
+        self.use_images = use_images
+        self.use_text_inventory = use_text_inventory
+        self.resolution = resolution  # low, medium, high
+
+        self.action_history = []
+        self.inventory_history = []
+        self.inventory_counters = []
+
+        self.tokens_used = 0
+
+        # set up dialogue history with few-shot prompt
+        self.set_up_few_shot_prompt()
+        self.system_prompt_dialogue = self.system_prompt()
+
+        self.dialogue_history = copy(self.prompt_examples)
+        self.images = copy(self.prompt_images)
+        self.initial_dialogue_length = len(self.dialogue_history)
+
+    def system_prompt(self):
+        # kept separate from dialogue history because certain models deal with system prompt differently
+        system_prompt_text = get_system_prompt(handlers=self.action_handlers)
+        if self.use_multimodal_content_format:
+            return {
+                "role": "system",
+                "content": [{"text": system_prompt_text, "type": "text"}],
+            }
+        return {
+            "role": "system",
+            "content": system_prompt_text,
+        }
+
+    def set_up_few_shot_prompt(self):
+        self.prompt_examples = []
+        self.prompt_images = []
+
+        if self.few_shot:
+            self.prompt_examples = get_prompt_example(
+                self.action_handlers,
+                use_text_inventory=self.use_text_inventory,
+                use_multimodal_content_format=self.use_multimodal_content_format,
+                use_images=self.use_images,
+            )
+            if self.use_images:
+                self.prompt_images = load_prompt_images(resolution=self.resolution)
+
+    def add_message_to_history(self, content: str | dict, role="user"):
+        if role == "assistant":
+            logger.info(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: SmeltAction | MoveAction):
+        if isinstance(action, SmeltAction) or isinstance(action, MoveAction):
+            self.action_history.append(action.model_dump())
+
+    def add_inventory_to_history(self, inventory: dict):
+        self.inventory_history.append(inventory)
+        # count inventory
+        counter = Counter()
+        for slot, item in inventory.items():
+            # ignore slot 0
+            if slot == 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 slot, item in observation["inventory"].items():
+            #     if item["quantity"] > 0:
+            #         clean_inv.append(item)
+            self.add_inventory_to_history(observation["inventory"])
+        if "image" in observation:
+            self.add_image_to_history(observation["image"])
+
+    def __str__(self):
+        return str(self.dialogue_history)
+
+    def reset(self):
+        # reset dialogue history to few-shot prompt
+        self.dialogue_history = copy(self.prompt_examples)
+        self.images = copy(self.prompt_images)
+        self.initial_dialogue_length = len(self.dialogue_history)
+
+        self.action_history = []
+        self.inventory_history = []
+        self.inventory_counters = []
+
+        self.tokens_used = 0
+
+    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),
+            "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 :]
+        )
 
 
 def get_downloaded_models() -> dict:
@@ -35,51 +207,3 @@ def get_torch_device() -> torch.device:
         else:
             device = torch.device("mps")
     return device
-
-
-def resize_image(img, target_resolution=(128, 128)):
-    if type(img) is np.ndarray:
-        img = cv2.resize(img, target_resolution, interpolation=cv2.INTER_LINEAR)
-    elif type(img) is torch.Tensor:
-        img = F.interpolate(img, size=target_resolution, mode="bilinear")
-    else:
-        raise ValueError("Unsupported type for img")
-    return img
-
-
-def save_frames_to_video(frames: list, out_path: str):
-    imgs = []
-    for id, (frame, goal) in enumerate(frames):
-        # if torch.is_tensor(frame):
-        # frame = frame.permute(0, 2, 3, 1).cpu().numpy()
-
-        frame = resize_image(frame, (320, 240)).astype("uint8")
-        cv2.putText(
-            frame,
-            f"FID: {id}",
-            (10, 25),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.8,
-            (255, 255, 255),
-            2,
-        )
-        cv2.putText(
-            frame,
-            f"Goal: {goal}",
-            (10, 55),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.8,
-            (255, 0, 0),
-            2,
-        )
-        imgs.append(Image.fromarray(frame))
-    imgs = imgs[::3]
-    imgs[0].save(
-        out_path,
-        save_all=True,
-        append_images=imgs[1:],
-        optimize=False,
-        quality=0,
-        duration=150,
-        loop=0,
-    )

plancraft-0.3.1.dist-info/METADATA

@@ -0,0 +1,253 @@
+Metadata-Version: 2.4
+Name: plancraft
+Version: 0.3.1
+Summary: Plancraft: an evaluation dataset for planning with LLM agents
+License: MIT License
+        
+        Copyright (c) 2024 Gautier Dagan
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Requires-Python: >=3.9
+Requires-Dist: hydra-core>=1.3.2
+Requires-Dist: imageio>=2.36.0
+Requires-Dist: loguru
+Requires-Dist: networkx
+Requires-Dist: numpy
+Requires-Dist: openai
+Requires-Dist: pandas
+Requires-Dist: pillow>=8.0.0
+Requires-Dist: pydantic
+Requires-Dist: pydantic-settings>=2.6.1
+Requires-Dist: scikit-learn>=1.6.0
+Requires-Dist: setuptools
+Requires-Dist: tqdm
+Requires-Dist: wandb
+Provides-Extra: full
+Requires-Dist: accelerate; extra == 'full'
+Requires-Dist: einops; extra == 'full'
+Requires-Dist: hf-transfer; extra == 'full'
+Requires-Dist: huggingface-hub; extra == 'full'
+Requires-Dist: matplotlib; extra == 'full'
+Requires-Dist: peft; extra == 'full'
+Requires-Dist: seaborn; extra == 'full'
+Requires-Dist: torch>=2.5.0; extra == 'full'
+Requires-Dist: torchvision>=0.20.0; extra == 'full'
+Requires-Dist: transformers>=4.43.3; extra == 'full'
+Description-Content-Type: text/markdown
+
+# plancraft
+
+[![Test](https://github.com/gautierdag/plancraft/actions/workflows/test.yaml/badge.svg)](https://github.com/gautierdag/plancraft/actions/workflows/test.yaml)
+![Python Version](https://img.shields.io/badge/python-3.9+-blue)
+![Ruff](https://img.shields.io/badge/linter-ruff-blue)
+[![PyPI Version](https://img.shields.io/pypi/v/plancraft)](https://pypi.org/project/plancraft/)
+[![Docker Pulls](https://img.shields.io/docker/pulls/gautierdag/plancraft)](https://hub.docker.com/r/gautierdag/plancraft)
+![License](https://img.shields.io/github/license/gautierdag/plancraft)
+![GitHub Repo stars](https://img.shields.io/github/stars/gautierdag/plancraft?style=social)
+
+[Paper](https://arxiv.org/abs/2412.21033) | [Website](https://gautierdag.github.io/plancraft/)
+
+Plancraft is a minecraft environment and agent that innovates on planning LLM agents with an oracle RAG retriever.
+
+You can install the package by running the following command:
+
+```bash
+pip install plancraft
+```
+
+![gif-example3](docs/images/train_images/TRAIN0010.gif)
+![gif-example1](docs/images/train_images/TRAIN1133.gif)
+![gif-example2](docs/images/train_images/TRAIN0383.gif)
+![gif-example3](docs/images/train_images/TRAIN1000.gif)
+
+The package provides a multimodal environment and dataset for evaluating planning agents. The dataset consists of examples of crafting tasks in Minecraft, where the agent must craft a target object from a set of initial items. The environment is a simplified version of Minecraft where the agent can move items between slots in an inventory and smelt items to create new items. The agent must craft the target object by moving or smelting items around in the inventory.
+
+## Usage
+
+The package provides a `PlancraftEnvironment` class that can be used to interact with the environment. Here is an example of how to use it:
+
+```python
+from plancraft.environments.env import PlancraftEnvironment
+
+
+def main():
+    # Create the environment with an inventory containing 10 iron ores and 23 oak logs
+    env = PlancraftEnvironment(
+        inventory={
+          10: dict(type="iron_ore", quantity=10),
+          23: dict(type="oak_log", quantity=23)
+        }
+    )
+    # move one log to slot 1
+    move_action = MoveAction(from_slot=23, to_slot=1, quantity=1)
+    observation = env.step(move_action)
+    # observation["inventory"] contains the updated symbolic inventory
+    # observation["image"] contains the updated image of the inventory
+
+    # smelt one iron ore
+    smelt_action = SmeltAction(from_slot=10, to_slot=11, quantity=1)
+    observation = env.step(smelt_action)
+
+    # no op
+    observation = env.step()
+```
+
+Note that the environment is deterministic and stateful, so the same action will always lead to the same observation and the environment will keep track of the state of the inventory.
+
+### Evaluator
+
+The package also provides an `Evaluator` class that can be used to evaluate the performance of an agent on our specific dataset. Here is an example of how to use it:
+
+```python
+from plancraft.evaluator import Evaluator
+from plancraft.config import EvalConfig
+
+def main():
+    # Create the config
+    config = EvalConfig(...)
+    # create model -- Note you can create your own model by subclassing PlancraftBaseModel
+    model = get_model(config)
+    # Create the evaluator
+    evaluator = Evaluator(config, model=model)
+    # Evaluate the agent
+    evaluator.eval_all_seeds()
+```
+
+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.
+
+#### The Evaluator interactive loop
+
+The evaluator loop for each example is as follows:
+
+```python
+# Initialize success and non-environment actions counter
+success = False
+num_non_env_actions = 0
+
+# Reset the environment and example
+reset(example)
+
+# Run the evaluation loop
+while not history.check_stuck() and history.num_steps < max_steps:
+    if action == StopAction:  # StopAction ends the episode
+        success = example.impossible  # Success if task is impossible
+        break
+    elif isinstance(action, str) and num_non_env_actions < 3:  
+        # Handle external tool action (str message)
+        observation = {"message": action}
+        num_non_env_actions += 1
+    else:  
+        # Handle environment action
+        if isinstance(action, str):  
+            # Handle invalid case (exceeded non-env action limit)
+            observation = environment.step()
+        else:
+            history.add_action_to_history(action)  # Add action to history
+            observation = environment.step(action)
+
+        # Convert observation to message and reset non-env counter
+        observation["target"] = example.target
+        observation["message"] = convert_observation_to_message(observation)
+        num_non_env_actions = 0
+
+        # Check if episode is complete
+        success = check_done(observation["inventory"], example.target)
+
+    # Update history with observation and message
+    history.add_observation_to_history(observation)
+    history.add_message_to_history(content=observation["message"], role="user")
+
+    if success:  # Exit loop if success
+        break
+
+    # Model predicts next action
+    raw_action = model.step(observation, dialogue_history=history)
+
+    # Update history with predicted action
+    history.add_message_to_history(content=raw_action, role="assistant")
+
+    # Parse raw action into a structured format
+    action = parse_raw_model_response(raw_action)
+
+# Return results after evaluation
+return {
+    "success": success,
+    "recipe_type": example.recipe_type,
+    "complexity": example.complexity,
+    "number_of_steps": history.num_steps,
+    "model_trace": history.trace(),
+    "example_id": example.id,
+    "impossible": example.impossible,
+}
+```
+
+### Observation
+
+The observation returned by the `PlancraftEnvironment` class is a dictionary with the following keys: `inventory` and `image`. The `inventory` key contains a dictionary with the slot number as the key and the item in the slot as the value (eg {"type": "iron_ingot", "quantity": 2}). The `image` key contains a numpy array representing the image of the inventory.
+
+The observation returned by the `Evaluator` class is a dictionary with the following keys: `inventory`, `image`, `message`, and `target`. The `message` key contains a string representing the environment formatted in text (we follow the annotation scheme described in our paper). The `target` key contains a string representing the target object to be crafted.
+
+### Implementing a Model
+
+To implement a model, you need to subclass the `PlancraftBaseModel` class and implement the `step` and `reset` method. See the `plancraft.models.dummy` module for an example of how to implement a basic model.
+
+You should then be able to use the `Evaluator` class to evaluate it.
+
+## Reproducing the Results tables in the paper
+
+To reproduce the results tables in the paper, you can use the `exps.sh` script in the `root` directory. The script will run the evaluation for all the models and seeds specified in the paper. The results will be saved in output directory but also on wandb if you have an account and set the `WANDB_API_KEY` environment variable.
+
+## Docker
+
+There is a docker image built to incorporate the latest code and its dependencies. I build it by running the following command:
+
+```bash
+docker buildx build --platform linux/amd64,linux/arm64 -t gautierdag/plancraft --push .
+```
+
+The image is available on [Docker Hub](https://hub.docker.com/r/gautierdag/plancraft). Note that, unlike the package, the docker image includes everything in the repo.
+
+## To Do
+
+Non-exhaustive list of things to do from highest to lowest priority:
+
+- [ ] Add minecraft wiki scrape and non-oracle search for pages
+- [ ] Improve planner to bring closer to optimal (the oracle planner does not consider  future crafting steps when moving items -- see paper for more details)
+- [ ] Rerun image models with better bounding box model
+  - [ ] Track bounding box accuracy
+- [ ] Implement a version of the image environment entirely on cuda/pytorch rather than cpu
+
+## PRs Welcomed
+
+If you would like to contribute to the project, please feel free to open a PR. I am happy to review and merge PRs that improve the project. If you have any questions, feel free to create an issue or reach out to me directly.
+
+## Citation
+
+```bibtex
+@misc{dagan2024plancraftevaluationdatasetplanning,
+      title={Plancraft: an evaluation dataset for planning with LLM agents}, 
+      author={Gautier Dagan and Frank Keller and Alex Lascarides},
+      year={2024},
+      eprint={2412.21033},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2412.21033}, 
+}
+```