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

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

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