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

plancraft/config.py

@@ -0,0 +1,155 @@
+from typing import Literal, Optional, Union
+
+from pydantic import BaseModel, model_validator
+
+try:
+    from plancraft.environments.recipes import RECIPES
+except ImportError:
+    RECIPES = {}
+
+DatasetSplit = Literal[
+    "train", "val", "val.small", "val.small.easy", "test", "test.small"
+]
+
+
+class EnvironmentConfig(BaseModel):
+    symbolic: bool
+    symbolic_observation_space: bool
+    symbolic_action_space: bool
+    preferred_spawn_biome: str = "plains"
+    resolution: list[int] = [512, 512]
+
+
+class PlancraftConfig(BaseModel):
+    model: str
+    adapter: str = ""
+    tokenizer: str
+    num_generations: int
+    mode: Literal["react", "act", "oracle", "dummy"] = "react"
+    output_dir: str
+    max_steps: int = 30  # max number of steps (smelt/move) to take in the environment before stopping
+    quantize: Literal[False, "int4", "int8"]
+    environment: EnvironmentConfig
+    split: DatasetSplit = "val.small"
+    max_message_window: int = 30  # max number of messages to keep in dialogue history (30 is around 8k llama3 tokens)
+    hot_cache: bool = True  # whether to cache the dialogue history between steps
+    resume: bool = True  # resume inference
+    few_shot: bool = True  # whether to use few-shot prompt
+    system_prompt: bool = True  # whether to use system prompt
+    valid_actions: list[str] = ["move", "smelt", "think", "search", "impossible"]
+    use_maskrcnn: bool = False  # whether to use maskrcnn for multimodal parsing
+
+    # observations
+    use_text_inventory: bool = True  # whether to include inventory in text
+    use_images: bool = False  # whether to include images in multimodal content
+    use_multimodal_content_format: bool = (
+        False  # whether to use multimodal content format
+    )
+
+    @model_validator(mode="after")
+    def validate(self):
+        assert set(
+            self.valid_actions
+        ).issubset(
+            {"move", "smelt", "think", "search", "impossible"}
+        ), "valid_actions should be subset of {'move', 'smelt', 'think', 'search', 'impossible'}"
+
+        if self.use_images:
+            assert (
+                not self.environment.symbolic
+            ), "Set environment.symbolic to False when using images"
+
+        return self
+
+
+class WandbConfig(BaseModel):
+    project: str
+    entity: str
+    mode: str
+
+
+class LaunchConfig(BaseModel):
+    command: str
+    job_name: str
+    gpu_limit: int
+    gpu_product: str
+    cpu_request: int
+    ram_request: str
+    interactive: bool = False
+    namespace: str = "informatics"
+    env_vars: dict[str, dict[str, str]]
+
+
+class EvalConfig(BaseModel):
+    plancraft: PlancraftConfig
+    wandb: WandbConfig
+    launch: LaunchConfig
+
+
+class TrainingArgs(BaseModel):
+    base_model: str = "llama3"
+    trace_mode: str = "oa"
+    push_to_hub: bool = False
+
+    # uses less space but not working with multi-gpu training..
+    qlora: bool = False
+
+    lora_alpha: int = 16
+    lora_dropout: float = 0.1
+    lora_r: int = 64
+    # training data args
+    seed: int = 42
+    # model args
+    batch_size: int = 1
+    max_seq_length: int = 8142
+    max_message_window: int = 100
+    only_assistant: bool = True
+
+    # training args
+    gradient_accumulation_steps: int = 4
+    learning_rate: float = 2e-4
+    max_grad_norm: float = 0.3
+    warmup_ratio: float = 0.03
+    num_train_epochs: int = 3
+    num_workers: int = 1
+
+
+class TrainConfig(BaseModel):
+    training: TrainingArgs
+    wandb: WandbConfig
+    launch: LaunchConfig
+
+
+class PlancraftExample(BaseModel):
+    target: str
+    inventory: dict[str, int]
+    slotted_inventory: list[dict[str, Union[str, int]]]
+    num_distractors: int
+    impossible: bool
+    optimal_path_length: Optional[int]
+    optimal_path: Optional[list[str]]
+    inventory_trace: Optional[list[dict[str, int]]]
+    items_used: Optional[int]
+    unique_items_used: Optional[int]
+    complexity: Optional[int]
+    complexity_bin: int
+    unseen_in_train: bool
+    unseen_in_val: bool
+    split: DatasetSplit
+    id: str
+
+    recipe_type: Optional[str] = ""
+
+    # post processing set recipe type
+    def model_post_init(self, __context):
+        recipe_types = set()
+        if self.optimal_path is None:
+            self.recipe_type = "impossible"
+            return
+        for step in self.optimal_path:
+            for r in RECIPES[step]:
+                recipe_types.add(r.recipe_type)
+        if len(recipe_types) == 1:
+            self.recipe_type = recipe_types.pop()
+        else:
+            self.recipe_type = "mixed"

plancraft/environments/actions.py

@@ -0,0 +1,218 @@
+from typing import Union
+
+from pydantic import BaseModel, field_validator, model_validator
+
+
+def convert_to_slot_index(slot: str) -> int:
+    slot = slot.strip()
+    grid_map = {
+        "[0]": 0,
+        "[A1]": 1,
+        "[A2]": 2,
+        "[A3]": 3,
+        "[B1]": 4,
+        "[B2]": 5,
+        "[B3]": 6,
+        "[C1]": 7,
+        "[C2]": 8,
+        "[C3]": 9,
+    }
+    if slot in grid_map:
+        return grid_map[slot]
+    else:
+        return int(slot[2:-1]) + 9
+
+
+def convert_from_slot_index(slot_index: int) -> str:
+    grid_map = {
+        0: "[0]",
+        1: "[A1]",
+        2: "[A2]",
+        3: "[A3]",
+        4: "[B1]",
+        5: "[B2]",
+        6: "[B3]",
+        7: "[C1]",
+        8: "[C2]",
+        9: "[C3]",
+    }
+    if slot_index < 10:
+        return grid_map[slot_index]
+    else:
+        return f"[I{slot_index-9}]"
+
+
+class SymbolicMoveAction(BaseModel):
+    """ "Moves an item from one slot to another"""
+
+    slot_from: int
+    slot_to: int
+    quantity: int
+    action_type: str = "move"
+
+    @field_validator("action_type", mode="before")
+    def fix_action_type(cls, value) -> str:
+        return "move"
+
+    @field_validator("slot_from", "slot_to", mode="before")
+    def transform_str_to_int(cls, value) -> int:
+        # if value is a string like [A1] or [I1], convert it to an integer
+        if isinstance(value, str):
+            try:
+                return convert_to_slot_index(value)
+            except ValueError:
+                raise AttributeError(
+                    "slot_from and slot_to must be [0] or [A1] to [C3] or [I1] to [I36]"
+                )
+        return value
+
+    @field_validator("quantity", mode="before")
+    def transform_quantity(cls, value) -> int:
+        if isinstance(value, str):
+            try:
+                return int(value)
+            except ValueError:
+                raise AttributeError("quantity must be an integer")
+        return value
+
+    @model_validator(mode="after")
+    def validate(self):
+        if self.slot_from == self.slot_to:
+            raise AttributeError("slot_from and slot_to must be different")
+        if self.slot_from < 0 or self.slot_from > 45:
+            raise AttributeError("slot_from must be between 0 and 45")
+        if self.slot_to < 1 or self.slot_to > 45:
+            raise AttributeError("slot_to must be between 1 and 45")
+        if self.quantity < 1 or self.quantity > 64:
+            raise AttributeError("quantity must be between 1 and 64")
+
+    def to_action_dict(self) -> dict:
+        return {
+            "inventory_command": [self.slot_from, self.slot_to, self.quantity],
+        }
+
+
+class SymbolicSmeltAction(BaseModel):
+    """Smelts an item and moves the result into a new slot"""
+
+    slot_from: int
+    slot_to: int
+    quantity: int
+    action_type: str = "smelt"
+
+    @field_validator("action_type", mode="before")
+    def fix_action_type(cls, value) -> str:
+        return "smelt"
+
+    @field_validator("slot_from", "slot_to", mode="before")
+    def transform_str_to_int(cls, value) -> int:
+        # if value is a string like [A1] or [I1], convert it to an integer
+        if isinstance(value, str):
+            try:
+                return convert_to_slot_index(value)
+            except ValueError:
+                raise AttributeError(
+                    "slot_from and slot_to must be [0] or [A1] to [C3] or [I1] to [I36]"
+                )
+        return value
+
+    @field_validator("quantity", mode="before")
+    def transform_quantity(cls, value) -> int:
+        if isinstance(value, str):
+            try:
+                return int(value)
+            except ValueError:
+                raise AttributeError("quantity must be an integer")
+        return value
+
+    @model_validator(mode="after")
+    def validate(self):
+        if self.slot_from == self.slot_to:
+            raise AttributeError("slot_from and slot_to must be different")
+        if self.slot_from < 0 or self.slot_from > 45:
+            raise AttributeError("slot_from must be between 0 and 45")
+        if self.slot_to < 1 or self.slot_to > 45:
+            raise AttributeError("slot_to must be between 1 and 45")
+        if self.quantity < 1 or self.quantity > 64:
+            raise AttributeError("quantity must be between 1 and 64")
+
+    def to_action_dict(self) -> dict:
+        return {
+            "smelt": [self.slot_from, self.slot_to, self.quantity],
+        }
+
+
+class ThinkAction(BaseModel):
+    """Think about the answer before answering"""
+
+    thought: str
+
+    def to_action_dict(self) -> dict:
+        return {}
+
+
+class SearchAction(BaseModel):
+    """Searches for a relevant document in the wiki"""
+
+    search_string: str
+
+    def to_action_dict(self) -> dict:
+        return {
+            "search": self.search_string,
+        }
+
+
+class RealActionInteraction(BaseModel):
+    mouse_direction_x: float = 0
+    mouse_direction_y: float = 0
+    right_click: bool = False
+    left_click: bool = False
+
+    @field_validator("mouse_direction_x", "mouse_direction_y")
+    def prevent_zero(cls, v):
+        if v > 10:
+            return 10
+        elif v < -10:
+            return -10
+        return v
+
+    def to_action_dict(self) -> dict:
+        return {
+            "camera": [self.mouse_direction_x, self.mouse_direction_y],
+            "use": int(self.right_click),
+            "attack": int(self.left_click),
+        }
+
+
+class StopAction(BaseModel):
+    """
+    Action that model can take to stop planning - decide impossible to continue
+    Note: also known as the "impossible" action
+    """
+
+    reason: str = ""
+
+
+class NoOp(SymbolicMoveAction):
+    """No operation action - special instance of move"""
+
+    def __init__(self):
+        super().__init__(slot_from=0, slot_to=1, quantity=1)
+        self.slot_to = 0
+
+    def __call__(self, *args, **kwargs):
+        return None
+
+    def __str__(self):
+        return "NoOp"
+
+
+# when symbolic action is true, can either move objects around or smelt
+SymbolicAction = SymbolicMoveAction  # | SymbolicSmeltAction
+
+# when symbolic action is false, then need to use mouse to move things around, but can use smelt action
+RealAction = RealActionInteraction | SymbolicSmeltAction
+
+
+class PydanticSymbolicAction(BaseModel):
+    root: Union[SymbolicMoveAction, SymbolicSmeltAction]

plancraft/environments/env_real.py

@@ -0,0 +1,316 @@
+from typing import Sequence, Union
+
+import numpy as np
+import json
+
+from loguru import logger
+
+from plancraft.environments.actions import RealAction
+
+try:
+    from minerl.env import _singleagent
+    from minerl.herobraine.env_specs.human_controls import HumanControlEnvSpec
+    from minerl.herobraine.hero import handlers, mc, spaces
+    from minerl.herobraine.hero.handler import Handler
+    from minerl.herobraine.hero.handlers.agent.action import Action
+    from minerl.herobraine.hero.handlers.agent.start import InventoryAgentStart
+    from minerl.herobraine.hero.handlers.translation import TranslationHandler
+
+    class InventoryCommandAction(Action):
+        """
+        Handler which lets agents programmatically interact with an open container
+
+        Using this - agents can move a chosen quantity of items from one slot to another.
+        """
+
+        def to_string(self):
+            return "inventory_command"
+
+        def xml_template(self) -> str:
+            return str("<InventoryCommands/>")
+
+        def __init__(self):
+            self._command = "inventory_command"
+            # first argument is the slot to take from
+            # second is the slot to put into
+            # third is the count to take
+            super().__init__(
+                self.command,
+                spaces.Tuple(
+                    (
+                        spaces.Discrete(46),
+                        spaces.Discrete(46),
+                        spaces.Discrete(64),
+                    )
+                ),
+            )
+
+        def from_universal(self, x):
+            return np.array([0, 0, 0], dtype=np.int32)
+
+    class SmeltCommandAction(Action):
+        """
+        An action handler for smelting an item
+        We assume smelting is immediate.
+        @TODO: might be interesting to explore using the smelting time as an additional planning parameter.
+
+        Using this agents can smelt items in their inventory.
+        """
+
+        def __init__(self):
+            self._command = "smelt"
+            # first argument is the slot to take from
+            # second is the slot to put into
+            # third is the count to smelt
+            super().__init__(
+                self.command,
+                spaces.Tuple(
+                    (
+                        spaces.Discrete(46),
+                        spaces.Discrete(46),
+                        spaces.Discrete(64),
+                    )
+                ),
+            )
+
+        def to_string(self):
+            return "smelt"
+
+        def xml_template(self) -> str:
+            return str("<SmeltCommands/>")
+
+        def from_universal(self, x):
+            return np.array([0, 0, 0], dtype=np.int32)
+
+    class InventoryResetAction(Action):
+        def __init__(self):
+            self._command = "inventory_reset"
+            super().__init__(self._command, spaces.Text([1]))
+
+        def to_string(self) -> str:
+            return "inventory_reset"
+
+        def to_hero(self, inventory_items: list[dict]):
+            return "{} {}".format(self._command, json.dumps(inventory_items))
+
+        def xml_template(self) -> str:
+            return "<InventoryResetCommands/>"
+
+        def from_universal(self, x):
+            return []
+
+    MINUTE = 20 * 60
+
+    class CustomInventoryAgentStart(InventoryAgentStart):
+        def __init__(self, inventory: list[dict[str, Union[str, int]]]):
+            super().__init__({item["slot"]: item for item in inventory})
+
+    class CraftingTableOnly(Handler):
+        def to_string(self):
+            return "start_with_crafting_table"
+
+        def xml_template(self) -> str:
+            return "<CraftingTableOnly>true</CraftingTableOnly>"
+
+    class InventoryObservation(TranslationHandler):
+        """
+        Handles GUI Workbench Observations for selected items
+        """
+
+        def to_string(self):
+            return "inventory"
+
+        def xml_template(self) -> str:
+            return str("""<ObservationFromFullInventory flat="false"/>""")
+
+        def __init__(self, item_list, _other="other"):
+            item_list = sorted(item_list)
+            super().__init__(
+                spaces.Dict(
+                    spaces={
+                        k: spaces.Box(
+                            low=0,
+                            high=2304,
+                            shape=(),
+                            dtype=np.int32,
+                            normalizer_scale="log",
+                        )
+                        for k in item_list
+                    }
+                )
+            )
+            self.num_items = len(item_list)
+            self.items = item_list
+
+        def add_to_mission_spec(self, mission_spec):
+            pass
+
+        def from_hero(self, info):
+            return info["inventory"]
+
+        def from_universal(self, obs):
+            raise NotImplementedError(
+                "from_universal not implemented in InventoryObservation"
+            )
+
+    class PlancraftBaseEnvSpec(HumanControlEnvSpec):
+        def __init__(
+            self,
+            symbolic_action_space=False,
+            symbolic_observation_space=False,
+            max_episode_steps=2 * MINUTE,
+            inventory: Sequence[dict] = (),
+            preferred_spawn_biome: str = "plains",
+            resolution=[260, 180],
+        ):
+            self.inventory = inventory
+            self.preferred_spawn_biome = preferred_spawn_biome
+            self.symbolic_action_space = symbolic_action_space
+            self.symbolic_observation_space = symbolic_observation_space
+
+            mode = "real"
+            if symbolic_action_space:
+                mode += "-symbolic-act"
+            else:
+                mode += "-real-act"
+
+            if symbolic_observation_space:
+                mode += "-symbolic-obs"
+
+            if symbolic_action_space:
+                cursor_size = 1
+            else:
+                cursor_size = 16
+
+            name = f"plancraft-{mode}-v0"
+            super().__init__(
+                name=name,
+                max_episode_steps=max_episode_steps,
+                resolution=resolution,
+                cursor_size_range=[cursor_size, cursor_size],
+            )
+
+        def create_agent_start(self) -> list[Handler]:
+            base_agent_start_handlers = super().create_agent_start()
+            return base_agent_start_handlers + [
+                CustomInventoryAgentStart(self.inventory),
+                handlers.PreferredSpawnBiome(self.preferred_spawn_biome),
+                handlers.DoneOnDeath(),
+                CraftingTableOnly(),
+            ]
+
+        def create_observables(self) -> list[TranslationHandler]:
+            if self.symbolic_observation_space:
+                return [
+                    handlers.POVObservation(self.resolution),
+                    InventoryObservation([item["slot"] for item in self.inventory]),
+                ]
+            return [handlers.POVObservation(self.resolution)]
+
+        def create_server_world_generators(self) -> list[Handler]:
+            # TODO the original biome forced is not implemented yet. Use this for now.
+            return [handlers.DefaultWorldGenerator(force_reset=True)]
+
+        def create_server_quit_producers(self) -> list[Handler]:
+            return [
+                handlers.ServerQuitFromTimeUp(
+                    (self.max_episode_steps * mc.MS_PER_STEP)
+                ),
+                handlers.ServerQuitWhenAnyAgentFinishes(),
+            ]
+
+        def create_server_initial_conditions(self) -> list[Handler]:
+            return [
+                handlers.TimeInitialCondition(allow_passage_of_time=False),
+                handlers.SpawningInitialCondition(allow_spawning=True),
+            ]
+
+        def create_actionables(self) -> list[TranslationHandler]:
+            """
+            Symbolic env can move items around in the inventory using function
+            Real env can use camera/keyboard
+            """
+            # Camera and mouse
+            if self.symbolic_action_space:
+                return [
+                    InventoryCommandAction(),
+                    SmeltCommandAction(),
+                    InventoryResetAction(),
+                ]
+            return [
+                handlers.KeybasedCommandAction(v, v) for k, v in mc.KEYMAP.items()
+            ] + [
+                handlers.CameraAction(),
+                SmeltCommandAction(),
+                InventoryResetAction(),
+            ]
+
+        def is_from_folder(self, folder: str) -> bool:
+            return False
+
+        def create_agent_handlers(self) -> list[Handler]:
+            return []
+
+        def create_mission_handlers(self):
+            return []
+
+        def create_monitors(self):
+            return []
+
+        def create_rewardables(self):
+            return []
+
+        def create_server_decorators(self) -> list[Handler]:
+            return []
+
+        def determine_success_from_rewards(self, rewards: list) -> bool:
+            return False
+
+        def get_docstring(self):
+            return self.__class__.__doc__
+
+    class RealPlancraft(_singleagent._SingleAgentEnv):
+        def __init__(
+            self,
+            inventory: list[dict],
+            preferred_spawn_biome="plains",
+            symbolic_action_space=False,
+            symbolic_observation_space=True,
+            resolution=[512, 512],
+            crop=True,
+        ):
+            # NOTE: crop is only supported for resolution 512x512 (default)
+            self.crop = crop
+            self.resolution = resolution
+            env_spec = PlancraftBaseEnvSpec(
+                symbolic_action_space=symbolic_action_space,
+                symbolic_observation_space=symbolic_observation_space,
+                preferred_spawn_biome=preferred_spawn_biome,
+                inventory=inventory,
+                resolution=resolution,
+            )
+            super(RealPlancraft, self).__init__(env_spec=env_spec)
+            self.reset()
+
+        def step(self, action: RealAction | dict):
+            if not isinstance(action, dict):
+                action = action.to_action_dict()
+            obs, rew, done, info = super().step(action)
+            if "pov" in obs and self.crop and self.resolution == [512, 512]:
+                # crop at position x=174, y=170 with width=164 and height=173
+                obs["pov"] = obs["pov"][174 : 174 + 164, 170 : 168 + 173]
+            return obs, rew, done, info
+
+        def fast_reset(self, new_inventory: list[dict]):
+            super().step({"inventory_reset": new_inventory})
+
+
+except ImportError:
+
+    class RealPlancraft:
+        def __init__(self, *args, **kwargs):
+            logger.warning(
+                "The 'minerl' package is required to use RealPlancraft. "
+                "Please install it using 'pip install plancraft[full]' or 'pip install minerl'."
+            )
+            raise ImportError("minerl package not found")