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

plancraft/models/oam.py

@@ -0,0 +1,283 @@
+from typing import Optional
+
+from loguru import logger
+
+import torch
+import torch.nn as nn
+import torchvision.transforms.v2 as v2
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedModel,
+)
+
+from plancraft.models.bbox_model import IntegratedBoundingBoxModel
+
+
+class PlancraftOAMConfig(PretrainedConfig):
+    model_type = "plancraft-aom"
+    is_composition = True
+
+    def __init__(
+        self,
+        from_llama=False,
+        **kwargs,
+    ):
+        self.from_llama = from_llama
+        super().__init__(**kwargs)
+
+
+class PlancraftOAM(PreTrainedModel):
+    config_class = PlancraftOAMConfig
+
+    def __init__(self, config: PlancraftOAMConfig):
+        super().__init__(config)
+
+        self.config = config
+        # load text model
+        if self.config.from_llama:
+            self.text_model = AutoModelForCausalLM.from_pretrained(
+                "meta-llama/Meta-Llama-3.1-8B-Instruct",
+            )
+        else:
+            text_model_config = AutoConfig.from_pretrained(
+                "meta-llama/Meta-Llama-3.1-8B-Instruct",
+            )
+            self.text_model = AutoModelForCausalLM.from_config(text_model_config)
+
+        # load vision model
+        self.vision_model = IntegratedBoundingBoxModel.from_pretrained(
+            "gautierdag/plancraft-maskrcnn"
+        )
+        self.vision_model.eval()
+
+        # convert vision features to text embedding
+        self.vision_to_text_embedding = nn.Linear(
+            1024, self.text_model.config.hidden_size
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            "meta-llama/Meta-Llama-3.1-8B-Instruct",
+            trust_remote=True,
+        )
+        # add special tokens
+        self.tokenizer.add_special_tokens(
+            {
+                "additional_special_tokens": [
+                    "<|inventory|>",
+                ]
+            }
+        )
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.inventory_idx = self.tokenizer.convert_tokens_to_ids("<|inventory|>")
+
+        # resize token embeddings
+        self.text_model.resize_token_embeddings(len(self.tokenizer))
+        # image transforms
+        self.transforms = v2.Compose(
+            [v2.ToImage(), v2.ToDtype(torch.float32, scale=True)]
+        )
+
+    @torch.no_grad()
+    def extract_bboxes(self, images: list) -> list[dict]:
+        if len(images) == 0:
+            return []
+        img_tensors = torch.stack([self.transforms(img) for img in images])
+        img_tensors = img_tensors.cuda()
+        # disable gradients
+        self.vision_model.freeze()
+        # get bounding box predictions
+        bbox_preds = self.vision_model(img_tensors)
+        return bbox_preds
+
+    def prepare_messages(self, messages: list, bboxes: list[dict]) -> str:
+        # no bounding boxes
+        if len(bboxes) == 0:
+            text = self.tokenizer.apply_chat_template(
+                messages, add_generation_prompt=not self.training, tokenize=False
+            )
+            text = text.replace("<|begin_of_text|>", "")
+            return text
+
+        # expand <|inventory|> tokens into N tokens (N = number of bounding boxes)
+        new_messages = []
+        i_pred = 0
+        for m in messages:
+            new_message = m.copy()
+            if new_message["role"] == "user" and new_message["content"].endswith(
+                "<|inventory|>"
+            ):
+                # add inventory tokens for each bounding box
+                new_message["content"] = new_message["content"].replace(
+                    "<|inventory|>",
+                    "<|inventory|>" * (bboxes[i_pred]["features"].shape[0]),
+                )
+                i_pred += 1
+            new_messages.append(new_message)
+
+        assert i_pred == len(
+            bboxes
+        ), "Number of inventory tokens does not match number of images"
+        # add special tokens
+
+        text = self.tokenizer.apply_chat_template(
+            new_messages, add_generation_prompt=not self.training, tokenize=False
+        )
+        text = text.replace("<|begin_of_text|>", "")
+        return text
+
+    def inputs_merger(
+        self,
+        input_ids: torch.LongTensor,
+        inputs_embeds: Optional[torch.Tensor],
+        image_hidden_states: Optional[torch.Tensor],
+    ):
+        # along batch dimension
+        for i in range(len(image_hidden_states)):
+            if len(image_hidden_states[i]) == 0:
+                assert (
+                    input_ids[i] == self.inventory_idx
+                ).sum() == 0, "No images but inventory token is still present"
+                continue
+
+            # count the number of inventory tokens
+            n_inventory_tokens = (input_ids[i] == self.inventory_idx).sum()
+            if n_inventory_tokens != image_hidden_states[i].shape[0]:
+                logger.warning(
+                    f"Number of inventory tokens ({n_inventory_tokens}) does not match number of bounding boxes ({image_hidden_states[i].shape[0]}). Possible truncation."
+                )
+                # truncated from the start
+                image_hidden_states[i] = image_hidden_states[i][-n_inventory_tokens:]
+
+            # replace inventory tokens with bbox features
+            inputs_embeds[i, input_ids[i] == self.inventory_idx] = image_hidden_states[
+                i
+            ]
+        return inputs_embeds
+
+    def process_inputs(
+        self,
+        batch_messages: list[list[dict]] = [],  # list of list of messages (untokenized)
+        batch_images: list[list] = [],  # list of list of images (unprocessed)
+    ) -> tuple[dict[str, torch.FloatTensor], list[torch.FloatTensor], int]:
+        """
+        Converts raw images and messages into model inputs
+        """
+        assert len(batch_images) == len(
+            batch_messages
+        ), "Number of images and messages should match in the batch dim"
+        # initial forward pass
+        texts_batch = []
+        image_hidden_states = []
+        total_boxes = 0
+        for images, messages in zip(batch_images, batch_messages):
+            # process images
+            bboxes = self.extract_bboxes(images)
+            if len(bboxes) > 0:
+                # get bbox features
+                features = torch.concat([p["features"] for p in bboxes], dim=0)
+                # upscale to text embedding size
+                features_embeds = self.vision_to_text_embedding(features)
+                image_hidden_states.append(features_embeds)
+                # count bboxes total
+                total_boxes += features.shape[0]
+            else:
+                image_hidden_states.append([])
+
+            # process messages
+            text = self.prepare_messages(messages, bboxes)
+            texts_batch.append(text)
+
+        # tokenize text
+        # @NOTE: truncation could cause issues with inventory tokens not matching number of boxes
+        # in that case, we will truncate the boxes from the end, and issue a warning
+        batch = self.tokenizer(
+            texts_batch,
+            truncation=True,
+            padding=True,
+            max_length=16384,
+            return_tensors="pt",
+        )
+        return batch, image_hidden_states, total_boxes
+
+    def forward(
+        self,
+        batch_messages: list[list[dict]] = [],  # list of list of messages (untokenized)
+        batch_images: list[list] = [],  # list of list of images (unprocessed)
+        **kwargs,
+    ):
+        labels = None
+        batch, image_hidden_states, total_boxes = self.process_inputs(
+            batch_messages, batch_images
+        )
+        # move to cuda
+        batch = {k: v.cuda() for k, v in batch.items()}
+        attention_mask = batch["attention_mask"]
+        input_ids = batch["input_ids"]
+
+        labels = input_ids.clone()
+        # remove inventory tokens from labels
+        labels[labels == self.inventory_idx] = -100
+        # sanity check: should have same number of boxes as inventory tokens
+        assert (labels == -100).sum() == total_boxes
+
+        # get text embeddings
+        inputs_embeds = self.text_model.get_input_embeddings()(input_ids)
+        inputs_embeds = self.inputs_merger(
+            input_ids, inputs_embeds, image_hidden_states
+        )
+        # forward pass
+        return self.text_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            labels=labels,
+            return_dict=True,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        batch_messages: list[list[dict]],
+        batch_images: list[list],
+        do_sample=True,
+        temperature=0.6,
+        max_new_tokens=32,
+    ):
+        self.tokenizer.padding_side = "left"
+
+        batch, image_hidden_states, _ = self.process_inputs(
+            batch_messages, batch_images
+        )
+        batch = {k: v.cuda() for k, v in batch.items()}
+        attention_mask = batch["attention_mask"]
+        input_ids = batch["input_ids"]
+
+        inputs_embeds = self.text_model.get_input_embeddings()(input_ids)
+        inputs_embeds = self.inputs_merger(
+            input_ids, inputs_embeds, image_hidden_states
+        )
+
+        generated_sequences = self.text_model.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            do_sample=do_sample,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            pad_token_id=self.tokenizer.eos_token_id,
+        )
+
+        # Decode the output
+        text_responses = self.tokenizer.batch_decode(
+            generated_sequences,
+            # generated_sequences[:, prompt_tokens:],
+            skip_special_tokens=False,
+        )
+
+        # remove <|eot_id|> tokens
+        text_responses = [
+            text_response.replace("<|eot_id|>", "") for text_response in text_responses
+        ]
+
+        _, total_tokens_used = generated_sequences.shape
+        return text_responses, total_tokens_used

plancraft/models/oracle.py

@@ -0,0 +1,265 @@
+import copy
+from collections import Counter
+
+from plancraft.config import EvalConfig
+from plancraft.environments.actions import (
+    RealActionInteraction,
+    StopAction,
+    SymbolicMoveAction,
+    SymbolicSmeltAction,
+)
+from plancraft.environments.planner import optimal_planner
+from plancraft.environments.recipes import (
+    ShapedRecipe,
+    ShapelessRecipe,
+    SmeltingRecipe,
+    id_to_item,
+)
+from plancraft.environments.sampler import MAX_STACK_SIZE
+from plancraft.models.base import ABCModel, History
+
+
+def item_set_id_to_type(item_set_ids: set[int]):
+    return set(id_to_item(i) for i in item_set_ids)
+
+
+def find_free_inventory_slot(inventory: list[dict], from_slot: int) -> int:
+    # find a free slot in the inventory for the item in from_slot
+    from_item_type, from_item_quantity = None, None
+
+    type_to_slot = {}
+    slot_to_quantity = {}
+    for item in inventory:
+        if ("slot" in item and item["slot"] == from_slot) or (
+            "index" in item and item["index"] == from_slot
+        ):
+            from_item_quantity = item["quantity"]
+            from_item_type = item["type"]
+            # break
+        item_type = item["type"]
+        quantity = item["quantity"]
+        if quantity == 0:
+            item_type = "air"
+
+        if "slot" in item:
+            item_slot = item["slot"]
+        else:
+            item_slot = item["index"]
+
+        if item_type not in type_to_slot:
+            type_to_slot[item_type] = [item_slot]
+        else:
+            type_to_slot[item_type].append(item_slot)
+
+        if item_slot not in slot_to_quantity:
+            slot_to_quantity[item_slot] = quantity
+        else:
+            slot_to_quantity[item_slot] += quantity
+
+    assert from_item_type is not None, f"Item not found in slot {from_slot}"
+
+    # if there is a free slot with the same item type
+    if from_item_type in type_to_slot:
+        for slot in type_to_slot[from_item_type]:
+            if (
+                slot != from_slot
+                and slot_to_quantity[slot] + from_item_quantity
+                <= MAX_STACK_SIZE[from_item_type]
+            ):
+                return slot
+
+    # if there is a free slot with air
+    for slot in type_to_slot["air"]:
+        if slot != from_slot and slot > 10:
+            return slot
+
+    raise ValueError("No free slot found")
+
+
+def find_item_in_inventory(target: str, inventory: list[dict]) -> int:
+    for item in inventory:
+        if item["type"] == target and item["quantity"] > 0:
+            if "slot" in item:
+                return item["slot"]
+            elif "index" in item:
+                return item["index"]
+            raise ValueError("Neither slot or index is set")
+
+
+def get_inventory_counter(inventory: list[dict]) -> Counter:
+    counter = Counter()
+    for item in inventory:
+        if "slot" in item and item["slot"] == 0:
+            continue
+        if "index" in item and item["index"] == 0:
+            continue
+        if item["type"] == "air":
+            continue
+        counter[item["type"]] += item["quantity"]
+    return counter
+
+
+def get_crafting_slot_item(inventory: list[dict]) -> dict:
+    for item in inventory:
+        if "slot" in item and item["slot"] == 0 and item["quantity"] > 0:
+            return item
+        if "index" in item and item["index"] == 0 and item["quantity"] > 0:
+            return item
+    return None
+
+
+def update_inventory(
+    inventory: list[dict], slot_from: int, slot_to: int, quantity: int
+) -> list[dict]:
+    """
+    decrements quantity of item in slot_from
+    NOTE: we don't care about incrementing the items in slot_to
+
+    """
+    new_inventory = []
+    for item in inventory:
+        if "slot" in item and item["slot"] == slot_from:
+            item["quantity"] -= quantity
+        elif "index" in item and item["index"] == slot_from:
+            item["quantity"] -= quantity
+        new_inventory.append(item)
+    return new_inventory
+
+
+class OracleModel(ABCModel):
+    """
+    Oracle model returns actions that solve the task optimally
+    """
+
+    def __init__(self, cfg: EvalConfig):
+        assert (
+            cfg.plancraft.environment.symbolic_action_space
+        ), "Only symbolic actions are supported for oracle"
+        self.history = History(objective="")
+        self.plans = []
+        self.subplans = []
+
+    def reset_history(self, objective: str = ""):
+        self.history.reset(objective=objective)
+        self.plans = []
+        self.subplans = []
+
+    def get_plan(self, observation: dict):
+        # objective="Craft an item of type: ...."
+        # this simply recovering the target item to craft
+        target = self.history.objective.split(": ")[-1]
+        inventory_counter = get_inventory_counter(observation["inventory"])
+        self.plans = optimal_planner(target=target, inventory=inventory_counter)
+
+    def get_next_action(
+        self, observation: dict
+    ) -> SymbolicMoveAction | SymbolicSmeltAction:
+        if len(self.subplans) > 0:
+            return self.subplans.pop(0)
+        if len(self.plans) == 0:
+            raise ValueError("No more steps in plan")
+
+        observed_inventory = copy.deepcopy(observation["inventory"])
+
+        # take item from crafting slot
+        if slot_item := get_crafting_slot_item(observed_inventory):
+            # move item from crafting slot to inventory
+            free_slot = find_free_inventory_slot(observed_inventory, from_slot=0)
+            return SymbolicMoveAction(
+                slot_from=0, slot_to=free_slot, quantity=slot_item["quantity"]
+            )
+
+        plan_recipe, new_inventory = self.plans.pop(0)
+        self.subplans = []
+        new_inventory_counter = Counter(new_inventory)
+        current_inventory = observed_inventory
+        current_inventory_counter = get_inventory_counter(current_inventory)
+        items_to_use_counter = current_inventory_counter - new_inventory_counter
+        new_items = new_inventory_counter - current_inventory_counter
+        assert len(new_items) == 1
+
+        if isinstance(plan_recipe, ShapelessRecipe):
+            crafting_slot = 1
+
+            # add each item to crafting slots
+            for item, quantity in items_to_use_counter.items():
+                n = 0
+                while n < quantity:
+                    from_slot = find_item_in_inventory(item, current_inventory)
+
+                    # skip if from_slot is the crafting slot
+                    if from_slot == crafting_slot:
+                        crafting_slot += 1
+                        n += 1
+                        continue
+
+                    # low_level_plan.append(("move", item, from_slot, crafting_slot, 1))
+                    action = SymbolicMoveAction(
+                        slot_from=from_slot, slot_to=crafting_slot, quantity=1
+                    )
+                    # update state of inventory
+                    current_inventory = update_inventory(
+                        current_inventory, from_slot, crafting_slot, 1
+                    )
+                    self.subplans.append(action)
+
+                    crafting_slot += 1
+                    n += 1
+
+        # if plan_recipe is a smelting recipe
+        elif isinstance(plan_recipe, SmeltingRecipe):
+            assert len(items_to_use_counter) == 1, "smelting only supports one item"
+            for item, quantity in items_to_use_counter.items():
+                from_slot = find_item_in_inventory(item, current_inventory)
+                free_slot = find_free_inventory_slot(
+                    current_inventory, from_slot=from_slot
+                )
+                action = SymbolicSmeltAction(
+                    slot_from=from_slot, slot_to=free_slot, quantity=quantity
+                )
+                self.subplans.append(action)
+
+        # if plan_recipe is a shaped recipe
+        elif isinstance(plan_recipe, ShapedRecipe):
+            for i, row in enumerate(plan_recipe.kernel):
+                for j, item_set in enumerate(row):
+                    inventory_position = (i * 3) + j + 1
+                    valid_items = item_set_id_to_type(item_set)
+                    for item in valid_items:
+                        if items_to_use_counter[item] > 0:
+                            from_slot = find_item_in_inventory(item, current_inventory)
+                            action = SymbolicMoveAction(
+                                slot_from=from_slot,
+                                slot_to=inventory_position,
+                                quantity=1,
+                            )
+                            items_to_use_counter[item] -= 1
+                            # update state of inventory
+                            current_inventory = update_inventory(
+                                current_inventory, from_slot, inventory_position, 1
+                            )
+                            self.subplans.append(action)
+                            break
+        else:
+            raise NotImplementedError(f"Recipe type {type(plan_recipe)} not supported")
+
+        return self.subplans.pop(0)
+
+    def step(
+        self, observation: dict
+    ) -> list[SymbolicMoveAction | RealActionInteraction | SymbolicSmeltAction]:
+        # add observation to history
+        self.history.add_observation_to_history(observation)
+
+        # get action
+        if len(self.plans) == 0:
+            self.get_plan(observation)
+            if self.plans is None:
+                self.plans = []
+                return StopAction()
+
+        action = self.get_next_action(observation)
+
+        # add action to history
+        self.history.add_action_to_history(action)
+        return action