From 41da76a058857b541327ba4634019a187005f831 Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Mon, 18 Nov 2024 00:23:14 +0000 Subject: [PATCH 01/15] Added a formatter class to simplify long conversational chains before going to the LLM to save tokens / context. --- src/datasetgen/mcts/conversation.py | 1 + src/datasetgen/mcts/conversation_formatter.py | 211 +++++++++++++----- src/tests/mcts/test_conversation_formatter.py | 139 ++++++++++++ 3 files changed, 299 insertions(+), 52 deletions(-) create mode 100644 src/tests/mcts/test_conversation_formatter.py diff --git a/src/datasetgen/mcts/conversation.py b/src/datasetgen/mcts/conversation.py index ea15a3786..252bf37e3 100644 --- a/src/datasetgen/mcts/conversation.py +++ b/src/datasetgen/mcts/conversation.py @@ -43,6 +43,7 @@ def entity_serializer(obj: Any) -> Dict: class Message(BaseModel): role: str content: str + metadata: Dict[str, Any] = field(default_factory=dict) class Conversation(BaseModel): diff --git a/src/datasetgen/mcts/conversation_formatter.py b/src/datasetgen/mcts/conversation_formatter.py index 5697e80d0..0aceaafe9 100644 --- a/src/datasetgen/mcts/conversation_formatter.py +++ b/src/datasetgen/mcts/conversation_formatter.py @@ -1,6 +1,7 @@ +import re from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import List, Dict, Any +from typing import List, Dict, Any, Tuple, Optional from datasetgen.mcts.conversation import Message, Conversation @@ -14,6 +15,110 @@ 4) Execution -- Taking into account 1,2 and 3, what steps do we need to take to successfully carry out the task """ + +class CodeProcessor: + """Handles code block processing and summarization""" + + @staticmethod + def extract_code_blocks(content: str) -> List[Tuple[str, int, int]]: + """ + Extract code blocks that are between comments. + Returns list of (code, start_line, end_line) tuples. + """ + lines = content.splitlines() + blocks = [] + current_block = [] + block_start = None + + for i, line in enumerate(lines, 1): # Start line counting from 1 + stripped = line.strip() + + if stripped.startswith('#'): + # If we have an active block, end it + if current_block: + blocks.append(( + '\n'.join(current_block), + block_start, + i - 1 + )) + current_block = [] + block_start = None + else: + # If this is first code line after comment + if stripped and block_start is None: + block_start = i + # Add non-empty lines to current block + if stripped: + current_block.append(line) + + # Handle any remaining block + if current_block: + blocks.append(( + '\n'.join(current_block), + block_start, + len(lines) + )) + + # If no blocks found, treat entire content as one block + if not blocks: + return [(content, 1, len(lines))] + + return blocks + + @staticmethod + def summarize_code_block(code: str, start_line: int = None, end_line: int = None, + preserve_comments: bool = True) -> str: + """ + Summarize a code block by replacing code sections with line count indicators. + Uses plural form when omitting multiple lines. + """ + lines = code.splitlines() + + # If there are no lines, return empty string + if not lines: + return "" + + # Check if this is a pure code block (no comments) + has_comments = any(line.strip().startswith('#') for line in lines) + + if not has_comments: + # For pure code blocks, treat the entire content as one section + return f"" + + # Regular processing for blocks with comments + result = [] + code_start = None + code_lines = 0 + + for i, line in enumerate(lines, 1): + stripped = line.strip() + + if stripped.startswith('#'): + # If we were collecting code lines + if code_start is not None: + # Determine whether to use singular or plural form + if code_lines == 1: + result.append(f"") + else: + result.append(f"") + code_start = None + code_lines = 0 + result.append(line) + else: + if stripped: + if code_start is None: + code_start = i + code_lines += 1 + + # Handle any remaining code section + if code_start is not None: + if code_lines == 1: + result.append(f"") + else: + result.append(f"") + + return '\n'.join(result) + class ConversationFormatter(ABC): """Abstract base class for conversation formatting strategies""" @@ -41,71 +146,73 @@ def format_conversation(self, conversation: Conversation) -> List[Message]: def format_message(self, message: Message) -> Message: return message -class OutputOnlyFormatter(ConversationFormatter): + +class StructurePreservingFormatter(ConversationFormatter): """ - Concrete formatter that preserves program outputs but omits program code. - Only includes the results/printed output from program execution. + Formatter that preserves program structure through comments while reducing token usage. """ - def __init__(self, planning = False): - self.planning = planning - - def format_conversation(self, conversation: Conversation) -> List[Message]: - formatted = [] - - # Keep system message unchanged - if conversation.messages and conversation.messages[0].role == "system": - formatted.append(self.format_message(conversation.messages[0])) - messages = conversation.messages[1:] - else: - messages = conversation.messages - - # Process message pairs (assistant's program + user's result) - for i in range(0, len(messages), 2): - if i + 1 >= len(messages): - break - - assistant_msg = messages[i] - user_msg = messages[i + 1] - - # Skip the assistant's program code - if assistant_msg.role != "assistant": - continue - - # Extract and format only the program output - formatted_result = self.format_message(user_msg) - if formatted_result: - formatted.append(formatted_result) + def __init__(self): + self.code_processor = CodeProcessor() - return formatted - - def format_message(self, message: Message) -> Message: + def format_message(self, message: Message, is_last: bool = False) -> Optional[Message]: if message.role == "system": - # Preserve system messages unchanged return Message(role="system", content=message.content) + elif message.role == "assistant": + if not is_last: # Summarize all but the last program + content = self.code_processor.summarize_code_block(message.content) + return Message( + role="assistant", + content=content, + metadata={"summarized": True} + ) + else: + return Message( + role="assistant", + content=message.content, + metadata={"summarized": False} + ) + elif message.role == "user": - # Extract only the program output/results content = message.content try: - # Split on common result markers - if "Execution result" in content: - content = content.split("Execution result")[1] - - # Remove state information but keep reward - if "Updated state:" in content: - content = content.split("Updated state:")[0] - - # Clean up and format - content = content.strip() - if content: + if "Execution result:" in content: + result = content.split("Execution result:")[1].split("Updated state:")[0] return Message( role="user", - content=f"{PLANNING_ADDITION_PROMPT if self.planning else ''} Previous attempt result: {content}".strip() + content=f"Execution result:\n{result.strip()}" + ) + else: + return Message( + role="user", + content=content.strip() ) except Exception as e: print(f"Error formatting user message: {str(e)}") return None - # Skip assistant messages (programs) - return None \ No newline at end of file + return None + + def format_conversation(self, conversation: Conversation) -> List[Message]: + formatted = [] + + # Handle system message if present + if conversation.messages and conversation.messages[0].role == "system": + formatted.append(self.format_message(conversation.messages[0])) + messages = conversation.messages[1:] + else: + messages = conversation.messages + + last_message_role = messages[-1].role + + # Format each message + for i, msg in enumerate(messages): + if last_message_role == 'assistant': + formatted_msg = self.format_message(msg, is_last=(i == len(messages) - 1)) + elif last_message_role == 'user': + formatted_msg = self.format_message(msg, is_last=(i == len(messages) - 2)) + if formatted_msg: + formatted.append(formatted_msg) + + return formatted \ No newline at end of file diff --git a/src/tests/mcts/test_conversation_formatter.py b/src/tests/mcts/test_conversation_formatter.py new file mode 100644 index 000000000..4d006441d --- /dev/null +++ b/src/tests/mcts/test_conversation_formatter.py @@ -0,0 +1,139 @@ +import unittest +from datasetgen.mcts.conversation import Message, Conversation + +from datasetgen.mcts.conversation_formatter import StructurePreservingFormatter, CodeProcessor + + +class TestStructurePreservingFormatter(unittest.TestCase): + def setUp(self): + self.formatter = StructurePreservingFormatter() + self.conversation = Conversation(messages=[ + Message( + role="system", + content="You are a helpful assistant." + ), + Message( + role="user", + content="Inventory: {}" + ), + Message( + role="assistant", + content="# Gather iron ore\nprint(0)\nprint(1)\n# Construct stone furnace" + ), + Message( + role="user", + content="Execution result: 1: 0\n2: 1" + ), + Message( + role="assistant", + content="# Gather more iron ore\nprint(2)\nprint(3)\n# Construct stone furnace" + ), + Message( + role="user", + content="Execution result: 3: 0\n4: 1" + ) + ]) + + def test_code_extractor(self): + code_snippets = CodeProcessor.extract_code_blocks(self.conversation.messages[2].content) + + # Should produce one code snippet + self.assertEqual(len(code_snippets), 1) + + def test_code_extractor_2(self): + code_snippets = CodeProcessor.extract_code_blocks(self.conversation.messages[2].content) + + self.assertIn("print(0)\nprint(1)", code_snippets[0]) + + def test_code_summariser(self): + code_block = "# Gather iron ore\nprint(0)\nprint(1)\n# Construct stone furnace" + code_snippets = CodeProcessor.extract_code_blocks(code_block) + + summarized = CodeProcessor.summarize_code_block(code_block, preserve_comments=True) + + self.assertEqual("# Gather iron ore\n\n# Construct stone furnace", summarized) + + summarized2 = CodeProcessor.summarize_code_block( + "# Gather iron ore\n# Gather more iron ore\nprint(0)\nprint(1)\n# Construct stone furnace", preserve_comments=True) + + self.assertEqual(summarized2, "# Gather iron ore\n# Gather more iron ore\n\n# Construct stone furnace") + + def test_format_conversation(self): + formatted = self.formatter.format_conversation(self.conversation) + + # Should produce 6 messages: system, user, assistant1, user1, assistant2, user2 + self.assertEqual(len(formatted), 6) + + # Check system message is preserved exactly + self.assertEqual(formatted[0].role, "system") + self.assertEqual(formatted[0].content, "You are a helpful assistant.") + + # Check first assistant message (not the last program, should be summarized) + assistant1 = formatted[2] + self.assertEqual(assistant1.role, "assistant") + expected_summary = ( + "# Gather iron ore\n" + "\n" + "# Construct stone furnace" + ) + self.assertEqual(assistant1.content, expected_summary) + + # Check last assistant message (should be preserved in full) + assistant2 = formatted[4] + self.assertEqual(assistant2.role, "assistant") + expected_full = ( + "# Gather more iron ore\n" + "print(2)\n" + "print(3)\n" + "# Construct stone furnace" + ) + self.assertEqual(assistant2.content, expected_full) + self.assertEqual(assistant2.metadata, {"summarized": False}) + + # Check user execution results are formatted correctly + user1 = formatted[3] + self.assertEqual(user1.role, "user") + self.assertEqual(user1.content, "Execution result:\n1: 0\n2: 1") + + user2 = formatted[5] + self.assertEqual(user2.role, "user") + self.assertEqual(user2.content, "Execution result:\n3: 0\n4: 1") + + def test_format_single_message(self): + # Test formatting of a non-last assistant message + message = Message( + role="assistant", + content="# First task\nprint(1)\n# Second task\nprint(2)\n" + ) + formatted = self.formatter.format_message(message, is_last=False) + expected = ( + "# First task\n" + "\n" + "# Second task\n" + "" + ) + self.assertEqual(formatted.content, expected) + self.assertEqual(formatted.metadata, {"summarized": True}) + + # Test formatting of a last assistant message (should preserve everything) + formatted_last = self.formatter.format_message(message, is_last=True) + self.assertEqual( + formatted_last.content, + "# First task\nprint(1)\n# Second task\nprint(2)\n" + ) + self.assertEqual(formatted_last.metadata, {"summarized": False}) + + def test_empty_code_blocks(self): + # Test handling of empty code blocks or blocks without comments + message = Message( + role="assistant", + content="print(1)\nprint(2)\n" + ) + formatted = self.formatter.format_message(message, is_last=False) + self.assertEqual( + formatted.content, + "" + ) + +if __name__ == '__main__': + unittest.main() \ No newline at end of file From 31fb2624f690368d04ff57d858755e5a90538f69 Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Mon, 18 Nov 2024 15:21:26 +0000 Subject: [PATCH 02/15] Instance logging using Rich. More data and some bug fixes. --- src/datasetgen/mcts/__main__.py | 28 ++- src/datasetgen/mcts/conversation.py | 10 +- src/datasetgen/mcts/conversation_formatter.py | 61 ++++-- src/datasetgen/mcts/db_client.py | 47 +++-- src/datasetgen/mcts/factorio_evaluator.py | 186 +++++++++++++----- src/datasetgen/mcts/logger.py | 135 +++++++++++++ src/datasetgen/mcts/mcts.py | 146 ++++++++++---- src/datasetgen/mcts/program.py | 3 + src/datasetgen/mcts/schema.sql | 2 +- src/llm_factory.py | 2 + src/tests/mcts/test_conversation_formatter.py | 30 ++- 11 files changed, 513 insertions(+), 137 deletions(-) create mode 100644 src/datasetgen/mcts/logger.py diff --git a/src/datasetgen/mcts/__main__.py b/src/datasetgen/mcts/__main__.py index 1c6bd50f4..c33fc060a 100644 --- a/src/datasetgen/mcts/__main__.py +++ b/src/datasetgen/mcts/__main__.py @@ -1,11 +1,16 @@ +import os +os.environ["FORCE_COLOR"] = "1" +os.environ["TERM"] = "xterm-256color" + import concurrent import os from typing import Tuple, List from dotenv import load_dotenv - +from rich import print +from rich.console import Console from cluster.local.cluster_ips import get_local_container_ips -from datasetgen.mcts.conversation_formatter import OutputOnlyFormatter +from datasetgen.mcts.conversation_formatter import StructurePreservingFormatter from datasetgen.mcts.db_client import DBClient from datasetgen.mcts.factorio_evaluator import FactorioEvaluator from datasetgen.mcts.game_state import GameState @@ -48,6 +53,10 @@ async def main(): password=os.getenv("SKILLS_DB_PASSWORD")) instances = create_parallel_instances() + + for instance in instances: + instance.speed(10) # Set the game speed to 10x normal speed for faster testing + # Initialize FactorioEvaluator with the list of instances evaluator = FactorioEvaluator(db_client, instances) initial_state = GameState.from_instance(instances[0]) @@ -59,12 +68,21 @@ async def main(): system_prompt = f.read().format(schema=instances[0].get_system_prompt()) print("Initializing MCTS...") - mcts = MCTS(llm, db_client, evaluator, system_prompt, initial_state, version=2, formatter=OutputOnlyFormatter(planning=True)) + console = Console(color_system="windows") + + mcts = MCTS(llm, + db_client, + evaluator, + system_prompt, + initial_state, + version=3, + version_description="Execution results exclude entities and inventory", + formatter=StructurePreservingFormatter(planning=True)) print("Starting MCTS search...") best_programs = await mcts.search( - n_iterations=100, - samples_per_iteration=5, + n_iterations=500, + samples_per_iteration=len(instances)-1, # One for each instance, minus a holdout. skip_failures=False, ) diff --git a/src/datasetgen/mcts/conversation.py b/src/datasetgen/mcts/conversation.py index 252bf37e3..d097fc4be 100644 --- a/src/datasetgen/mcts/conversation.py +++ b/src/datasetgen/mcts/conversation.py @@ -59,10 +59,8 @@ def parse_raw(cls, data: Dict[str, Any]) -> 'Conversation': def add_result(self, program: str, reward: float, response: str, new_state: GameState, entities: List[Union[Entity, EntityGroup]]): """Add program execution result to conversation""" self.messages.append(Message(role="assistant",content=program)) - self.messages.append(Message(role="user", content= - f"""Execution result (reward: {reward}): - {response} + self.messages.append(Message(role="user", content=f"Execution result: \n{response}")) - Updated state: - Inventory: {json.dumps(new_state.inventory.__dict__, indent=2)} - Entities: {json.dumps(entities, indent=2, cls=EntityEncoder)}""")) \ No newline at end of file + # Updated state: + # Inventory: {json.dumps(new_state.inventory.__dict__, indent=2)} + # Entities: {json.dumps(entities, indent=2, cls=EntityEncoder)}""")) \ No newline at end of file diff --git a/src/datasetgen/mcts/conversation_formatter.py b/src/datasetgen/mcts/conversation_formatter.py index 0aceaafe9..039169cb1 100644 --- a/src/datasetgen/mcts/conversation_formatter.py +++ b/src/datasetgen/mcts/conversation_formatter.py @@ -7,18 +7,31 @@ PLANNING_ADDITION_PROMPT = \ """ -First bring out a thorough step-by-step plan how you can achieve this task and then create the python script to achieve the task. +Identify an appropriate next task given the recent history and environment feedback. +Develop thorough step-by-step plan how you can achieve this task and then create the python script to achieve the task. For your plan, follow this structure: 1) What entities are needed for the task 2) What entities do we have on the map, in different entity inventories or in our inventory 3) What entities are we missing for the task -4) Execution -- Taking into account 1,2 and 3, what steps do we need to take to successfully carry out the task +4) Execution -- Taking into account 1,2 and 3, what are the steps we need to take to successfully carry out the task. """ class CodeProcessor: """Handles code block processing and summarization""" + @staticmethod + def is_comment_start(line: str) -> bool: + """Check if line starts a comment (either # or \"\"\")""" + stripped = line.strip() + return stripped.startswith('#') or stripped.startswith('"""') + + @staticmethod + def is_comment_end(line: str) -> bool: + """Check if line ends a comment block""" + stripped = line.strip() + return not stripped.startswith('"""') and stripped.endswith('"""') + @staticmethod def extract_code_blocks(content: str) -> List[Tuple[str, int, int]]: """ @@ -70,33 +83,44 @@ def summarize_code_block(code: str, start_line: int = None, end_line: int = None preserve_comments: bool = True) -> str: """ Summarize a code block by replacing code sections with line count indicators. - Uses plural form when omitting multiple lines. + Handles both inline (#) and block (\"\"\") comments. """ lines = code.splitlines() - # If there are no lines, return empty string if not lines: return "" - # Check if this is a pure code block (no comments) - has_comments = any(line.strip().startswith('#') for line in lines) - - if not has_comments: - # For pure code blocks, treat the entire content as one section - return f"" - - # Regular processing for blocks with comments result = [] code_start = None code_lines = 0 + in_docstring = False for i, line in enumerate(lines, 1): stripped = line.strip() + # Handle docstring boundaries + if stripped.startswith('"""'): + if not in_docstring: # Start of docstring + if code_start is not None: + if code_lines == 1: + result.append(f"") + else: + result.append(f"") + code_start = None + code_lines = 0 + in_docstring = True + result.append(line) + continue + + if in_docstring: + result.append(line) + if stripped.endswith('"""'): + in_docstring = False + continue + + # Handle regular comments and code if stripped.startswith('#'): - # If we were collecting code lines if code_start is not None: - # Determine whether to use singular or plural form if code_lines == 1: result.append(f"") else: @@ -150,10 +174,12 @@ def format_message(self, message: Message) -> Message: class StructurePreservingFormatter(ConversationFormatter): """ Formatter that preserves program structure through comments while reducing token usage. + It replaces all code not in the most recent history with `` tags. """ - def __init__(self): + def __init__(self, planning=True): self.code_processor = CodeProcessor() + self.planning = planning def format_message(self, message: Message, is_last: bool = False) -> Optional[Message]: if message.role == "system": @@ -179,9 +205,12 @@ def format_message(self, message: Message, is_last: bool = False) -> Optional[Me try: if "Execution result:" in content: result = content.split("Execution result:")[1].split("Updated state:")[0] + content = f"Execution result:\n{result.strip()}" + if self.planning: + content = PLANNING_ADDITION_PROMPT + '\n' + content return Message( role="user", - content=f"Execution result:\n{result.strip()}" + content=content ) else: return Message( diff --git a/src/datasetgen/mcts/db_client.py b/src/datasetgen/mcts/db_client.py index bbb06d0a2..7fdb1653e 100644 --- a/src/datasetgen/mcts/db_client.py +++ b/src/datasetgen/mcts/db_client.py @@ -14,28 +14,33 @@ def __init__(self, **db_config): self.conn = psycopg2.connect(**db_config) async def create_program(self, program: Program) -> Program: - with self.conn.cursor() as cur: - cur.execute(""" - INSERT INTO programs (code, value, visits, parent_id, state_json, conversation_json, completion_token_usage, prompt_token_usage, token_usage, response, holdout_value, raw_reward, version) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - RETURNING id, created_at - """, (program.code, program.value, program.visits, program.parent_id, - program.state.to_raw() if program.state else None, - json.dumps(program.conversation.dict()), - program.completion_token_usage, - program.prompt_token_usage, - program.token_usage, - program.response, - program.holdout_value, - program.raw_reward, - program.version - )) + try: + with self.conn.cursor() as cur: + cur.execute(""" + INSERT INTO programs (code, value, visits, parent_id, state_json, conversation_json, completion_token_usage, prompt_token_usage, token_usage, response, holdout_value, raw_reward, version, version_description) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + RETURNING id, created_at + """, (program.code, program.value, program.visits, program.parent_id, + program.state.to_raw() if program.state else None, + json.dumps(program.conversation.dict()), + program.completion_token_usage, + program.prompt_token_usage, + program.token_usage, + program.response, + program.holdout_value, + program.raw_reward, + program.version, + program.version_description + )) - id, created_at = cur.fetchone() - self.conn.commit() - program.id = id - program.created_at = created_at - return program + id, created_at = cur.fetchone() + self.conn.commit() + program.id = id + program.created_at = created_at + return program + except Exception as e: + print(e) + raise e @tenacity.retry(retry=retry_if_exception_type((psycopg2.OperationalError, psycopg2.InterfaceError)), wait=wait_exponential(multiplier=1, min=4, max=10)) async def sample_parent(self, version=1) -> Optional[Program]: diff --git a/src/datasetgen/mcts/factorio_evaluator.py b/src/datasetgen/mcts/factorio_evaluator.py index 5e7ce156b..00af76e87 100644 --- a/src/datasetgen/mcts/factorio_evaluator.py +++ b/src/datasetgen/mcts/factorio_evaluator.py @@ -1,62 +1,160 @@ import asyncio -from copy import deepcopy -from typing import List, Tuple, Optional, Union +from typing import List, Tuple, Union from datasetgen.mcts.db_client import DBClient from datasetgen.mcts.game_state import GameState +from datasetgen.mcts.logger import FactorioLogger from datasetgen.mcts.program import Program from factorio_entities import Entity, EntityGroup from factorio_instance import FactorioInstance class FactorioEvaluator: - def __init__(self, db_client: DBClient, instances: List[FactorioInstance]): + def __init__(self, db_client: DBClient, instances: List[FactorioInstance], value_accrual_time=10): self.db = db_client self.instances = instances[:-1] # Main instances self.holdout = instances[-1] # Holdout instance + self.value_accrual_time = value_accrual_time # Time to accrue value before evaluating + self.logger = FactorioLogger(len(instances)) + self.logger.start() + + def set_sampling_status(self): + for i in range(len(self.instances)): + self.logger.update_instance(i, status="sampling") async def evaluate_batch(self, programs: List[Program], start_state: GameState) -> List[Program]: - # Reset holdout to same starting state - self.holdout.reset(start_state) - holdout_future = asyncio.create_task(self._run_holdout()) - - # Evaluate programs from same starting state - eval_futures = [] - for prog, inst in zip(programs, self.instances): - inst.reset(start_state) # All instances start from same state - eval_futures.append(self._evaluate_single(prog, inst)) - - # Wait for all evaluations - eval_results = await asyncio.gather(*eval_futures) - holdout_value = await holdout_future - - print(f"Holdout value: {holdout_value}") - - # Update programs with relative rewards - for program, (raw_reward, state, response, entities) in zip(programs, eval_results): - relative_reward = raw_reward - holdout_value - print(f"Program {program.id}: raw={raw_reward}, relative={relative_reward}") - - # First update program state and conversation - program.value = relative_reward - program.state = state - program.raw_reward = raw_reward - program.holdout_value = holdout_value - program.conversation.add_result(program.code, relative_reward, response, state, entities) - program.response = response - - # Filter out programs with no state - these are - # Return updated programs - return programs - - async def _evaluate_single(self, program: Program, instance: FactorioInstance) -> Tuple[float, GameState, str, List[Union[Entity, EntityGroup]]]: - reward, _, result = instance.eval(program.code, timeout=60) - state = GameState.from_instance(instance) - entities = instance.get_entities() - return reward, state, result, entities + try: + # Reset holdout to same starting state + self.holdout.reset(start_state) + self.logger.update_instance(len(self.instances), status="running", program_id=None) + holdout_future = asyncio.create_task(self._run_holdout()) + + # Evaluate programs from same starting state + eval_futures = [] + for i, (prog, inst) in enumerate(zip(programs, self.instances)): + inst.reset(start_state) # All instances start from same state + self.logger.update_instance(i, program_id=prog.id, status="resetting") + eval_futures.append(self._evaluate_single(i, prog, inst)) + + # Wait for all evaluations + eval_results = await asyncio.gather(*eval_futures) + holdout_value = await holdout_future + + self.logger.update_instance(len(self.instances), + status="completed", + current_reward=holdout_value) + + # Update programs with relative rewards + for i, (program, (raw_reward, state, response, entities)) in enumerate(zip(programs, eval_results)): + relative_reward = raw_reward - holdout_value + + self.logger.update_instance( + i, + status="completed", + raw_reward=raw_reward, + holdout_value=holdout_value, + relative_reward=relative_reward, + total_programs=self.logger.instances[i].total_programs + 1, + ) + + # First update program state and conversation + program.value = relative_reward + program.state = state + program.raw_reward = raw_reward + program.holdout_value = holdout_value + program.conversation.add_result(program.code, relative_reward, response, state, entities) + program.response = response + + # Return updated programs + return programs + except Exception as e: + for i in range(len(self.instances) + 1): + self.logger.update_instance( + i, + status="error", + error_count=self.logger.instances[i].error_count + 1 + ) + raise e + + async def _evaluate_single(self, instance_id: int, program: Program, instance: FactorioInstance) -> Tuple[ + float, GameState, str, List[Union[Entity, EntityGroup]]]: + try: + start_entities = instance.get_entities() + start_inventory = instance.inspect_inventory() + self.logger.update_instance(instance_id, status="starting value") + initial_value, _ = instance.score() + + self.logger.update_instance(instance_id, status="executing") + reward, _, result = instance.eval(program.code, timeout=60) + + self.logger.update_instance(instance_id, status="capturing state") + state = GameState.from_instance(instance) + entities = instance.get_entities() + final_inventory = instance.inspect_inventory() + + self.logger.update_instance(instance_id, status="accruing value") + await asyncio.sleep(10) + + score, _ = instance.score() + final_reward = score - initial_value + + self.logger.update_instance( + instance_id, + status="accrued value", + current_reward=final_reward, + raw_reward=final_reward, + final_entities=len(entities), + start_entities=len(start_entities), + start_inventory_count=sum([v for k, v in start_inventory.__dict__.items() if v > 0]), + final_inventory_count=sum([v for k, v in final_inventory.__dict__.items() if v > 0]) + ) + + if "error" in result.lower(): + self.logger.update_instance( + instance_id, + status="error", + error_count=self.logger.instances[instance_id].error_count + 1 + ) + + return final_reward, state, result, entities + + except Exception as e: + self.logger.update_instance( + instance_id, + status="error", + error_count=self.logger.instances[instance_id].error_count + 1 + ) + raise e async def _run_holdout(self) -> float: """Run holdout instance for same duration as programs""" - await asyncio.sleep(10) # Same timeout as program evaluation - reward, _ = self.holdout.score() - return reward \ No newline at end of file + try: + initial_entities = self.holdout.get_entities() + start_inventory = self.holdout.inspect_inventory() + + initial_value, _ = self.holdout.score() + self.logger.update_instance(len(self.instances), status="accruing value") + await asyncio.sleep(10) + entities = self.holdout.get_entities() + reward, _ = self.holdout.score() + final_inventory = self.holdout.inspect_inventory() + + self.logger.update_instance(len(self.instances), + status="accrued value", + final_entities=len(entities), + start_entities=len(initial_entities), + start_inventory_count=sum([v for k, v in start_inventory.__dict__.items() if v > 0]), + final_inventory_count=sum([v for k, v in final_inventory.__dict__.items() if v > 0])) + + return reward - initial_value + except Exception as e: + self.logger.update_instance( + len(self.instances), + status="error", + error_count=self.logger.instances[len(self.instances)].error_count + 1 + ) + raise e + + def __del__(self): + """Clean up logger on deletion""" + self.logger.stop() \ No newline at end of file diff --git a/src/datasetgen/mcts/logger.py b/src/datasetgen/mcts/logger.py new file mode 100644 index 000000000..68ae6c401 --- /dev/null +++ b/src/datasetgen/mcts/logger.py @@ -0,0 +1,135 @@ +from dataclasses import dataclass, field +from datetime import datetime +from typing import Dict, List, Optional + +from rich.console import Console +from rich.layout import Layout +from rich.live import Live +from rich.panel import Panel +from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn +from rich.table import Table +from rich.text import Text + + +@dataclass +class InstanceMetrics: + """Tracks metrics for a single Factorio instance""" + instance_id: int + port: int + program_id: Optional[int] = None + current_reward: float = 0.0 + raw_reward: float = 0.0 + holdout_value: float = 0.0 + relative_reward: float = 0.0 + status: str = "idle" + last_update: datetime = field(default_factory=datetime.now) + error_count: int = 0 + total_programs: int = 0 + + start_entities: int = 0 + final_entities: int = 0 + + start_inventory_count: int = 0 + final_inventory_count: int = 0 + + version: int = 0 + version_description: str = "N/A" + + +class FactorioLogger: + def __init__(self, num_instances: int): + self.console = Console() + self.layout = Layout() + self.instances: Dict[int, InstanceMetrics] = {} + self.live: Optional[Live] = None + self.start_time = datetime.now() + + self.progress = Progress( + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TaskProgressColumn(), + expand=True + ) + self.progress_task = None + + # Initialize metrics for each instance + for i in range(num_instances): + self.instances[i] = InstanceMetrics(instance_id=i, port=34197 + i) + + def start(self): + """Start the live display with reduced refresh rate and transient=False""" + self.live = Live( + self._generate_layout(), + refresh_per_second=4, # Reduced from 4 + transient=False # Prevents clearing/redrawing + ) + self.live.start() + + def stop(self): + """Stop the live display""" + if self.live: + self.live.stop() + + def update_progress(self, advance: int = 1): + if self.progress_task is not None: + self.progress.update(self.progress_task, advance=advance) + if self.live: + self.live.update(self._generate_layout()) + + def update_instance(self, instance_id: int, **updates): + """Update metrics for a specific instance""" + if instance_id in self.instances: + instance = self.instances[instance_id] + for key, value in updates.items(): + if hasattr(instance, key): + setattr(instance, key, value) + instance.last_update = datetime.now() + if self.live: + self.live.update(self._generate_layout()) + + def _generate_instance_panel(self, instance: InstanceMetrics) -> Panel: + """Generate a panel for a single instance""" + table = Table(show_header=False, box=None, padding=(0, 1)) + avg_time = "N/A" + if instance.total_programs > 0: + avg_time = f"{((datetime.now() - self.start_time).total_seconds() / instance.total_programs):.2f} sec" + + # Add metrics rows + table.add_row("Program ID:", str(instance.program_id or "None")) + table.add_row("Status:", Text(instance.status, style="green" if instance.status == "running" else "yellow")) + table.add_row("Action Reward:", f"{instance.current_reward:.2f}") + table.add_row("Holdout Reward:", f"{instance.holdout_value:.2f}") + table.add_row("Advantage:", f"{instance.relative_reward:.2f}") + table.add_row("Error Count:", Text(str(instance.error_count), style="red" if instance.error_count > 0 else "white")) + table.add_row("Total Programs:", str(instance.total_programs)) + table.add_row("Last Update:", instance.last_update.strftime("%H:%M:%S")) + table.add_row("Avg Time / Program:", Text(avg_time)) + table.add_row("# Entities:", Text(f"{instance.start_entities} -> {instance.final_entities}", style="cyan" if instance.start_entities >= instance.final_entities else "red")) + table.add_row("# Inventory:", Text(f"{instance.start_inventory_count} -> {instance.final_inventory_count}")) + + return Panel( + table, + title=f"Instance {instance.instance_id} (Port: {instance.port})" if instance.instance_id != len(self.instances)-1 else f"Holdout Instance (Port: {instance.port})", + border_style="blue" + ) + + def _generate_layout(self) -> Layout: + layout = Layout() + + # Create instance metrics layout + metrics_layout = Layout() + metrics_layout.split_row( + *[Layout(self._generate_instance_panel(instance), ratio=1) + for instance in self.instances.values()] + ) + + # Add progress bar if it exists + if self.progress_task is not None: + layout.split_column( + Layout(metrics_layout, ratio=4), + Layout(self.progress, ratio=1) + ) + else: + layout = metrics_layout + + return layout \ No newline at end of file diff --git a/src/datasetgen/mcts/mcts.py b/src/datasetgen/mcts/mcts.py index 09467dc9d..98d75b507 100644 --- a/src/datasetgen/mcts/mcts.py +++ b/src/datasetgen/mcts/mcts.py @@ -2,6 +2,8 @@ import json from typing import Optional, List +from tenacity import wait_exponential, retry + from datasetgen.mcts.conversation import Conversation, Message from datasetgen.mcts.conversation_formatter import ConversationFormatter, DefaultFormatter, PLANNING_ADDITION_PROMPT from datasetgen.mcts.db_client import DBClient @@ -19,6 +21,7 @@ def __init__(self, initial_state: GameState, formatter: ConversationFormatter = DefaultFormatter(), version=1, + version_description="" ): self.llm = llm_factory @@ -27,60 +30,112 @@ def __init__(self, self.system_prompt = system_prompt self.initial_state = initial_state self.version = version + self.version_description=version_description self.formatter = formatter - def _extract_code(self, response) -> str: - content = response.content if "claude" in self.llm.model else response.choices[0].message.content - try: - content = content.split('```python')[1].split('```')[0].strip() - except IndexError as e: - pass + def _verify_response_is_python(self, content): + code = content # Parse into an AST to verify that this is a program try: - ast = compile(content, filename="", mode="exec") - except SyntaxError as e: + ast = compile(code, filename="", mode="exec") + except SyntaxError: # Take the last line off and try again - content = content.rsplit('\n', 1)[0] + '\n' - try: - ast = compile(content, filename="", mode="exec") - except SyntaxError as e1: - raise ValueError(f"Invalid Python code: {content}") from e - - return content.strip() + code = code.rsplit('\n', 1)[0] + '\n' + ast = compile(code, filename="", mode="exec") + return code + def _extract_code_from_choice(self, choice) -> Optional[str]: + """Extract code from a single completion choice""" + code = "" + try: + content = choice.message.content + code = self._verify_response_is_python(content) + return code + except Exception as e: + try: + content = choice.message.content + code = content.split('```python')[1].split('```')[0].strip() + code = self._verify_response_is_python(code) + return code + except Exception as e1: + print(f"Failed to extract code from choice: {str(e1)}") + # Sometimes it samples a leading line, before writing unblocked python code. + content = "\n".join(choice.message.content.split("\n")[1:]) + try: + code = self._verify_response_is_python(content) + return code + except Exception as e2: + print(f"Failed to extract code from choice after removing leading line: {str(e2)}") + + try: + code = content.split('from factorio_instance import *')[1].strip() + code = self._verify_response_is_python(code) + return code + except Exception as e2: + print(f"Failed to extract code from choice after removing leading line and factorio_instance import: {str(e2)}") + return None + + @retry(wait=wait_exponential(multiplier=1, min=4, max=10)) async def _generate_programs_batch(self, conversation: Conversation, n_samples: int) -> List[Program]: - tasks = [] - messages = conversation.model_dump()['messages'] + """Generate multiple programs in a single API call using 'n' parameter""" formatted_messages = self.formatter.to_llm_messages( self.formatter.format_conversation(conversation) ) + if "claude" in self.llm.model: + assert n_samples == 1, "Number of samples must be 1 for Claude" - async def generate_single(): - try: - response = self.llm.call(messages=formatted_messages) - code = self._extract_code(response) - completion_token_usage = response.usage.completion_tokens if hasattr(response, 'usage') else None - prompt_token_usage = response.usage.prompt_tokens if hasattr(response, 'usage') else None - token_usage = response.usage.total_tokens if hasattr(response, 'usage') else None - id_ = hash((code, json.dumps(messages))) # create an ID by hashing the code and conversation - return Program( - id=id_, - code=code, - conversation=conversation, - token_usage=token_usage, - completion_token_usage=completion_token_usage, - prompt_token_usage=prompt_token_usage, - version=self.version - ) - except Exception as e: - print(f"Program generation failed: {str(e)}") - return None - - for _ in range(n_samples): - tasks.append(generate_single()) - - return [p for p in await asyncio.gather(*tasks) if p is not None] + try: + # Single API call to generate n_samples completions + response = self.llm.call( + messages=formatted_messages, + n=n_samples, + temperature=0.7 # Adjust as needed + ) + + programs = [] + messages = conversation.model_dump()['messages'] + + # Process all choices from the response + if "claude" in self.llm.model: + + # Handle Claude response format + code = self._extract_code_from_choice(response) + if code: + programs.append(Program( + id=hash((code, json.dumps(messages))), + code=code, + conversation=conversation, + token_usage=response.usage.total_tokens if hasattr(response, 'usage') else None, + completion_token_usage=response.usage.completion_tokens if hasattr(response, 'usage') else None, + prompt_token_usage=response.usage.prompt_tokens if hasattr(response, 'usage') else None, + version=self.version, + version_description=self.version_description + )) + else: + # Handle OpenAI response format with multiple choices + for choice in response.choices: + code = self._extract_code_from_choice(choice) + if code: + programs.append(Program( + id=hash((code, json.dumps(messages))), + code=code, + conversation=conversation, + token_usage=response.usage.total_tokens // n_samples if hasattr(response, + 'usage') else None, + completion_token_usage=response.usage.completion_tokens // n_samples if hasattr(response, + 'usage') else None, + prompt_token_usage=response.usage.prompt_tokens // n_samples if hasattr(response, + 'usage') else None, + version=self.version, + version_description=self.version_description + )) + + return programs + + except Exception as e: + print(f"Batch program generation failed: {str(e)}") + return [] async def search(self, n_iterations: int, @@ -112,9 +167,11 @@ async def search(self, content=f"Inventory: {json.dumps(start_state.inventory.__dict__)}\n\n{PLANNING_ADDITION_PROMPT}") ]) + self.evaluator.set_sampling_status() # Generate multiple programs from same parent programs = await self._generate_programs_batch(conversation, samples_per_iteration) if not programs: + self.evaluator.logger.update_progress() continue # Filter out None programs @@ -141,4 +198,7 @@ async def search(self, except Exception as e: print(f"Evaluation failed for batch: {str(e)}") - continue \ No newline at end of file + self.evaluator.logger.update_progress() + continue + + self.evaluator.logger.update_progress() \ No newline at end of file diff --git a/src/datasetgen/mcts/program.py b/src/datasetgen/mcts/program.py index aa16c59b7..7bfeb7b64 100644 --- a/src/datasetgen/mcts/program.py +++ b/src/datasetgen/mcts/program.py @@ -25,6 +25,7 @@ class Program(BaseModel): token_usage: Optional[int] = None response: Optional[str] = None version: int = 1 + version_description: str = "" def get_uct(self, parent_visits: int, exploration_constant: float = 1.41) -> float: if self.visits == 0: @@ -54,4 +55,6 @@ def from_row(cls, row: Dict): completion_token_usage=row['completion_token_usage'], token_usage=row['token_usage'], response=row['response'], + version=row['version'], + version_description=row['version_description'] ) \ No newline at end of file diff --git a/src/datasetgen/mcts/schema.sql b/src/datasetgen/mcts/schema.sql index 766458a10..ad9002813 100644 --- a/src/datasetgen/mcts/schema.sql +++ b/src/datasetgen/mcts/schema.sql @@ -33,7 +33,7 @@ BEGIN FROM programs WHERE version = target_version ORDER BY created_at DESC - LIMIT 20 + LIMIT 100 ), softmax AS ( SELECT diff --git a/src/llm_factory.py b/src/llm_factory.py index f1da2542a..ad557d836 100644 --- a/src/llm_factory.py +++ b/src/llm_factory.py @@ -36,6 +36,7 @@ def remove_whitespace_blocks(self, messages): def call(self, *args, **kwargs): max_tokens = kwargs.get('max_tokens', 1500) model_to_use = kwargs.get('model', self.model) + n_samples = kwargs.get('n', 1) if "claude" in model_to_use: # Set up and call the Anthropic API api_key = os.getenv('ANTHROPIC_API_KEY') @@ -109,6 +110,7 @@ def call(self, *args, **kwargs): max_tokens = kwargs.get('max_tokens', 2048), temperature=kwargs.get('temperature', 0.3), messages=kwargs.get('messages', None), + n=n_samples, # Use requested number of samples #stop=["\n\n"],#, "\n#"], #presence_penalty=1, #frequency_penalty=0.6, diff --git a/src/tests/mcts/test_conversation_formatter.py b/src/tests/mcts/test_conversation_formatter.py index 4d006441d..2aff35637 100644 --- a/src/tests/mcts/test_conversation_formatter.py +++ b/src/tests/mcts/test_conversation_formatter.py @@ -6,7 +6,7 @@ class TestStructurePreservingFormatter(unittest.TestCase): def setUp(self): - self.formatter = StructurePreservingFormatter() + self.formatter = StructurePreservingFormatter(planning=False) self.conversation = Conversation(messages=[ Message( role="system", @@ -135,5 +135,33 @@ def test_empty_code_blocks(self): "" ) + def test_docstring_code_summariser(self): + code_block = '''from factorio_instance import * + + """ + Objective: We need to get 20 copper plates + + Planning: + 1. Print the recipe for copper plates + 2. Analyze the current game state + """ + + """ + Step 1: Print recipe for copper plates + """ + print("Copper Plate Recipe:") + print("Crafting requires smelting") + + """ + Step 2: Analyze current game state + """ + inventory = inspect_inventory() + print(f"Current inventory: {inventory}")''' + + summarized = self.formatter.code_processor.summarize_code_block(code_block, preserve_comments=True) + + pass + + if __name__ == '__main__': unittest.main() \ No newline at end of file From 713ef3a88baebef5b99d023c664f0f3ac58836dd Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Mon, 18 Nov 2024 19:30:57 +0000 Subject: [PATCH 03/15] Partial implementation of a chunked MCTS searcher --- src/datasetgen/mcts/__main__.py | 10 +- src/datasetgen/mcts/chunked_mcts.py | 139 ++++++++++++++++++++++++++++ src/tests/mcts/test_mcts_chunker.py | 129 ++++++++++++++++++++++++++ 3 files changed, 274 insertions(+), 4 deletions(-) create mode 100644 src/datasetgen/mcts/chunked_mcts.py create mode 100644 src/tests/mcts/test_mcts_chunker.py diff --git a/src/datasetgen/mcts/__main__.py b/src/datasetgen/mcts/__main__.py index c33fc060a..efbf58185 100644 --- a/src/datasetgen/mcts/__main__.py +++ b/src/datasetgen/mcts/__main__.py @@ -1,4 +1,7 @@ import os + +from datasetgen.mcts.chunked_mcts import ChunkedMCTS + os.environ["FORCE_COLOR"] = "1" os.environ["TERM"] = "xterm-256color" @@ -68,15 +71,14 @@ async def main(): system_prompt = f.read().format(schema=instances[0].get_system_prompt()) print("Initializing MCTS...") - console = Console(color_system="windows") - mcts = MCTS(llm, + mcts = ChunkedMCTS(llm, db_client, evaluator, system_prompt, initial_state, - version=3, - version_description="Execution results exclude entities and inventory", + version=4, + version_description="Step-wise evaluation / Execution results exclude entities and inventory", formatter=StructurePreservingFormatter(planning=True)) print("Starting MCTS search...") diff --git a/src/datasetgen/mcts/chunked_mcts.py b/src/datasetgen/mcts/chunked_mcts.py new file mode 100644 index 000000000..0bd99b6d5 --- /dev/null +++ b/src/datasetgen/mcts/chunked_mcts.py @@ -0,0 +1,139 @@ +import ast +import asyncio +import json +from dataclasses import dataclass +from typing import List, Tuple, Optional + +from datasetgen.mcts.conversation import Conversation, Message +from datasetgen.mcts.conversation_formatter import PLANNING_ADDITION_PROMPT +from datasetgen.mcts.game_state import GameState +from datasetgen.mcts.mcts import MCTS +from datasetgen.mcts.program import Program + + +class ChunkedMCTS(MCTS): + def _split_into_chunks(self, program_code: str) -> List[Program]: + chunks = [] + current_docstring = "" + current_code = [] + module = ast.parse(program_code) + + for node in module.body: + if isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant): + if current_code and current_docstring: + chunks.append(Program( + code=f'"""{current_docstring}"""\n'+"\n".join(current_code), + conversation=Conversation(messages=[]) # We need this to conform to the Program class structure + )) + current_code = [] + current_docstring = node.value.s + else: + # Get the source lines for this node + start = node.lineno + end = node.end_lineno if hasattr(node, 'end_lineno') else start + code_lines = program_code.splitlines()[start - 1:end] + current_code.extend(code_lines) + + if current_code and current_docstring: + chunks.append(Program( + code=f'"""{current_docstring}"""\n'+"\n".join(current_code), + conversation=Conversation(messages=[]) # We need this to conform to the Program class structure + )) + return chunks + + async def _evaluate_chunks(self, chunks: List[Program], start_state: GameState, instance_id: int) -> Tuple[ + List[Program], float]: + current_state = start_state + holdout_start = GameState.from_instance(self.evaluator.holdout) + self.evaluator.holdout.reset(holdout_start) + holdout_future = asyncio.create_task(self.evaluator._run_holdout()) + entity_list = [] + try: + for chunk in chunks: + instance = self.evaluator.instances[instance_id] + instance.reset(current_state) + self.evaluator.logger.update_instance(instance_id, status="executing") + + reward, state, response, entities = await self.evaluator._evaluate_single( + instance_id, + chunk, + instance + ) + entity_list.append(entities) + chunk.state = state + chunk.value = reward + chunk.response = response + + current_state = state + except Exception as e: + print(f"Error during evaluation: {e}") + raise e # Propagate the exception to handle it elsewhere if needed. + + holdout_value = await holdout_future + return chunks, holdout_value, entity_list + + async def search(self, n_iterations: int, samples_per_iteration: int, skip_failures: bool = False): + for iteration in range(n_iterations): + parent = await self.db.sample_parent(version=self.version) + start_state = parent.state if parent else self.initial_state + conversation = parent.conversation if parent else Conversation(messages=[ + Message(role="system", content=self.system_prompt), + Message(role="user", + content=f"Inventory: {json.dumps(start_state.inventory.__dict__)}\n\n{PLANNING_ADDITION_PROMPT}") + ]) + + self.evaluator.set_sampling_status() + raw_programs = await self._generate_programs_batch(conversation, samples_per_iteration) + + for i, (program, chunks) in enumerate(raw_programs): + instance_id = i % (len(self.evaluator.instances) - 1) + + try: + evaluated_chunks, holdout, entity_list = await self._evaluate_chunks(chunks, start_state, instance_id) + last_chunk_id = parent.id if parent else None + last_conversation_stage = program.conversation + for chunk, entities in zip(evaluated_chunks, entity_list): + chunk_program = Program( + code=chunk.code, + conversation=program.conversation, + value=chunk.value - (holdout / len(evaluated_chunks)), # Distribute holdout value + state=chunk.state, + response=chunk.response, + version=self.version, + version_description=self.version_description, + parent_id=last_chunk_id, + token_usage=program.token_usage // len(evaluated_chunks), + completion_token_usage=program.completion_token_usage // len(evaluated_chunks), + prompt_token_usage=program.prompt_token_usage // len(evaluated_chunks) + ) + last_conversation_stage.add_result(chunk.code, chunk.value - (holdout / len(evaluated_chunks)), chunk.response, chunk.state, entities) + chunk_program.id = hash((chunk.code, json.dumps(chunk_program.conversation.model_dump()['messages']))) + chunk_program.conversation = last_conversation_stage + + last_chunk_id = chunk_program.id + if not skip_failures or chunk_program.value is not None: + await self.db.create_program(chunk_program) + + except Exception as e: + print(f"Failed to evaluate program: {str(e)}") + continue + + self.evaluator.logger.update_progress() + + async def _generate_programs_batch(self, conversation: Conversation, n_samples: int) -> List[Tuple[Program, List[Program]]]: + programs = await super()._generate_programs_batch(conversation, n_samples) + chunked_programs = [] + + for i, program in enumerate(programs): + try: + chunks = self._split_into_chunks(program.code) + if not chunks: + continue + + chunked_programs.append((program, chunks)) + + except Exception as e: + print(f"Failed to process chunks for program: {str(e)}") + continue + + return chunked_programs \ No newline at end of file diff --git a/src/tests/mcts/test_mcts_chunker.py b/src/tests/mcts/test_mcts_chunker.py new file mode 100644 index 000000000..350a90f25 --- /dev/null +++ b/src/tests/mcts/test_mcts_chunker.py @@ -0,0 +1,129 @@ +import unittest +import ast +from dataclasses import dataclass +from typing import List, Optional + +from datasetgen.mcts.chunked_mcts import ChunkedMCTS + + +@dataclass +class ProgramChunk: + docstring: str + code: str + state: Optional['GameState'] = None + reward: float = 0.0 + + +class TestProgramChunkSplitter(unittest.TestCase): + def setUp(self): + self.splitter = ChunkedMCTS(None, None, None, "", None) + + def test_single_chunk(self): + code = ''' +"""First task""" +x = 1 +y = 2 +''' + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0].docstring, "First task") + self.assertEqual(chunks[0].code.strip(), "x = 1\ny = 2") + + def test_multiple_chunks(self): + code = ''' +"""First task""" +x = 1 + +"""Second task""" +y = 2 + +"""Third task""" +z = 3 +''' + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 3) + self.assertEqual([c.docstring for c in chunks], ["First task", "Second task", "Third task"]) + self.assertEqual([c.code.strip() for c in chunks], ["x = 1", "y = 2", "z = 3"]) + + def test_multiline_code(self): + code = ''' +"""Build structure""" +x = 1 +y = 2 +z = x + y + +"""Place items""" +items = [1, 2, 3] +for item in items: + place(item) +''' + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 2) + self.assertTrue("x = 1" in chunks[0].code) + self.assertTrue("for item in items:" in chunks[1].code) + + def test_multiline_docstring(self): + code = ''' +""" +Build initial structure +With multiple lines +""" +x = 1 + +"""Place items""" +y = 2 +''' + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 2) + self.assertEqual(chunks[0].docstring.strip(), "Build initial structure\nWith multiple lines") + + def test_no_docstring(self): + code = "x = 1\ny = 2" + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 0) + + def test_function_definitions(self): + code = ''' +"""Setup functions""" +def func1(): + return 1 + +def func2(): + x = 1 + return x + +"""Use functions""" +result = func1() + func2() +''' + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 2) + self.assertTrue("def func1():" in chunks[0].code) + self.assertTrue("def func2():" in chunks[0].code) + self.assertTrue("result = func1() + func2()" in chunks[1].code) + + def test_empty_chunks(self): + code = ''' +"""Task 1""" + +"""Task 2""" +x = 1 +''' + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0].docstring, "Task 2") + + def test_code_with_strings(self): + code = ''' +"""First task""" +x = "This is a string" +y = """This is a multiline +string but not a docstring""" +''' + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 1) + self.assertTrue('x = "This is a string"' in chunks[0].code) + self.assertTrue('y = """This is a multiline' in chunks[0].code) + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file From c946b8a87e1d30fe0c7de4c7066b35774b64fb5d Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Mon, 18 Nov 2024 20:44:10 +0000 Subject: [PATCH 04/15] Retain comments and whitespace in chunk parser --- src/datasetgen/mcts/chunked_mcts.py | 71 +++++++++++++------ src/datasetgen/mcts/game_state.py | 55 ++++++++++++-- src/tests/mcts/test_mcts_chunker.py | 29 ++++++-- .../mcts/test_save_load_python_namespace.py | 38 ++++++++++ 4 files changed, 159 insertions(+), 34 deletions(-) create mode 100644 src/tests/mcts/test_save_load_python_namespace.py diff --git a/src/datasetgen/mcts/chunked_mcts.py b/src/datasetgen/mcts/chunked_mcts.py index 0bd99b6d5..3c9119434 100644 --- a/src/datasetgen/mcts/chunked_mcts.py +++ b/src/datasetgen/mcts/chunked_mcts.py @@ -14,35 +14,58 @@ class ChunkedMCTS(MCTS): def _split_into_chunks(self, program_code: str) -> List[Program]: chunks = [] + lines = program_code.splitlines() current_docstring = "" - current_code = [] + current_lines = [] + current_start = 0 module = ast.parse(program_code) - for node in module.body: + # Track docstring positions to identify chunk boundaries + docstring_positions = [] + for i, node in enumerate(module.body): if isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant): - if current_code and current_docstring: - chunks.append(Program( - code=f'"""{current_docstring}"""\n'+"\n".join(current_code), - conversation=Conversation(messages=[]) # We need this to conform to the Program class structure - )) - current_code = [] - current_docstring = node.value.s + docstring_positions.append((node.lineno - 1, node.end_lineno - 1, node.value.s)) + + if not docstring_positions: + return [] + + # Process chunks between docstrings + for i, (start_pos, end_pos, docstring) in enumerate(docstring_positions): + # If this isn't the first chunk, append the previous one + if current_docstring and current_lines: + chunks.append(Program( + code="\n".join(current_lines), + conversation=Conversation(messages=[]) + )) + + # Start new chunk + current_docstring = docstring + current_lines = [] + + # Add any lines before the docstring that follow the previous chunk + if i > 0: + prev_end = docstring_positions[i - 1][1] + current_lines.extend(lines[prev_end + 1:start_pos]) + + # Add the docstring lines + current_lines.extend(lines[start_pos:end_pos + 1]) + + # If this is the last docstring, add all remaining lines + if i == len(docstring_positions) - 1: + current_lines.extend(lines[end_pos + 1:]) + chunks.append(Program( + code="\n".join(current_lines), + conversation=Conversation(messages=[]) + )) else: - # Get the source lines for this node - start = node.lineno - end = node.end_lineno if hasattr(node, 'end_lineno') else start - code_lines = program_code.splitlines()[start - 1:end] - current_code.extend(code_lines) - - if current_code and current_docstring: - chunks.append(Program( - code=f'"""{current_docstring}"""\n'+"\n".join(current_code), - conversation=Conversation(messages=[]) # We need this to conform to the Program class structure - )) + # Add lines until the next docstring + next_start = docstring_positions[i + 1][0] + current_lines.extend(lines[end_pos + 1:next_start]) + return chunks - async def _evaluate_chunks(self, chunks: List[Program], start_state: GameState, instance_id: int) -> Tuple[ - List[Program], float]: + async def _evaluate_chunks(self, chunks: List[Program], start_state: GameState, instance_id: int) \ + -> Tuple[List[Program], float]: current_state = start_state holdout_start = GameState.from_instance(self.evaluator.holdout) self.evaluator.holdout.reset(holdout_start) @@ -88,6 +111,10 @@ async def search(self, n_iterations: int, samples_per_iteration: int, skip_failu for i, (program, chunks) in enumerate(raw_programs): instance_id = i % (len(self.evaluator.instances) - 1) + + self.evaluator.instances[instance_id].reset(start_state) # Reset to the start state before evaluating each chunk. + self.logger.update_instance(i, program_id=program.id, status="resetting") + try: evaluated_chunks, holdout, entity_list = await self._evaluate_chunks(chunks, start_state, instance_id) last_chunk_id = parent.id if parent else None diff --git a/src/datasetgen/mcts/game_state.py b/src/datasetgen/mcts/game_state.py index 4066212ac..ac8709f3c 100644 --- a/src/datasetgen/mcts/game_state.py +++ b/src/datasetgen/mcts/game_state.py @@ -1,44 +1,89 @@ import json +import pickle import time from dataclasses import dataclass, field from enum import Enum from typing import List, Dict, Any +def is_serializable(obj: Any) -> bool: + """Test if an object can be serialized with pickle""" + try: + pickle.dumps(obj) + return True + except (pickle.PicklingError, TypeError, AttributeError): + return False + +def filter_serializable_vars(vars_dict: Dict[str, Any]) -> Dict[str, Any]: + """Filter dictionary to only include serializable items""" + return { + key: value for key, value in vars_dict.items() + if is_serializable(value) + } + @dataclass class GameState: """Serializable Factorio game state""" entities: str # Serialized list of entities inventory: Dict[str, int] timestamp: float = field(default_factory=time.time) + namespace: bytes = field(default_factory=bytes) @classmethod def from_instance(cls, instance: 'FactorioInstance') -> 'GameState': """Capture current game state from Factorio instance""" entities = instance._save_entity_state(compress=True, encode=True) + + # Filter and pickle only serializable variables + if hasattr(instance, 'persistent_vars'): + serializable_vars = filter_serializable_vars(instance.persistent_vars) + namespace = pickle.dumps(serializable_vars) if serializable_vars else bytes() + else: + namespace = bytes() + return cls( entities=entities, inventory=instance.inspect_inventory(), + namespace=namespace ) @classmethod def parse_raw(cls, json_str: str) -> 'GameState': data = json.loads(json_str) - return cls(entities=data['entities'], inventory=data['inventory'], timestamp=data['timestamp'] if 'timestamp' in data else time.time()) + namespace = bytes.fromhex(data['namespace']) if 'namespace' in data else bytes() + return cls( + entities=data['entities'], + inventory=data['inventory'], + timestamp=data['timestamp'] if 'timestamp' in data else time.time(), + namespace=namespace + ) @classmethod def parse(cls, data) -> 'GameState': - return cls(entities=data['entities'], inventory=data['inventory'], - timestamp=data['timestamp'] if 'timestamp' in data else time.time()) + namespace = bytes.fromhex(data['namespace']) if 'namespace' in data else bytes() + return cls( + entities=data['entities'], + inventory=data['inventory'], + timestamp=data['timestamp'] if 'timestamp' in data else time.time(), + namespace=namespace + ) def to_raw(self) -> str: """Convert state to JSON string""" return json.dumps({ 'entities': self.entities, 'inventory': self.inventory.__dict__, - 'timestamp': self.timestamp + 'timestamp': self.timestamp, + 'namespace': self.namespace.hex() if self.namespace else '' }) def to_instance(self, instance: 'FactorioInstance'): """Restore game state to Factorio instance""" instance._load_entity_state(self.entities, decode=True) - instance.set_inventory(**self.inventory) \ No newline at end of file + instance.set_inventory(**self.inventory) + + # Merge pickled namespace with existing persistent_vars + if self.namespace: + restored_vars = pickle.loads(self.namespace) + if not hasattr(instance, 'persistent_vars') or instance.persistent_vars is None: + instance.persistent_vars = {} + instance.persistent_vars.update(restored_vars) \ No newline at end of file diff --git a/src/tests/mcts/test_mcts_chunker.py b/src/tests/mcts/test_mcts_chunker.py index 350a90f25..8b726a536 100644 --- a/src/tests/mcts/test_mcts_chunker.py +++ b/src/tests/mcts/test_mcts_chunker.py @@ -26,8 +26,21 @@ def test_single_chunk(self): ''' chunks = self.splitter._split_into_chunks(code) self.assertEqual(len(chunks), 1) - self.assertEqual(chunks[0].docstring, "First task") - self.assertEqual(chunks[0].code.strip(), "x = 1\ny = 2") + self.assertTrue("First task" in chunks[0].code) + self.assertTrue( "x = 1\ny = 2" in chunks[0].code.strip()) + + def test_comment(self): + code = ''' +"""First task""" +# This is a comment +x = 1 +y = 2 +''' + chunks = self.splitter._split_into_chunks(code) + self.assertEqual(len(chunks), 1) + self.assertTrue("First task" in chunks[0].code) + self.assertTrue("x = 1\ny = 2" in chunks[0].code.strip()) + self.assertTrue("# This is a comment" in chunks[0].code.strip()) def test_multiple_chunks(self): code = ''' @@ -42,8 +55,10 @@ def test_multiple_chunks(self): ''' chunks = self.splitter._split_into_chunks(code) self.assertEqual(len(chunks), 3) - self.assertEqual([c.docstring for c in chunks], ["First task", "Second task", "Third task"]) - self.assertEqual([c.code.strip() for c in chunks], ["x = 1", "y = 2", "z = 3"]) + for program, task, value in zip(chunks, ["First task", "Second task", "Third task"], ["x = 1", "y = 2", "z = 3"]): + self.assertTrue(task in program.code) + self.assertTrue(value in program.code) + def test_multiline_code(self): code = ''' @@ -75,7 +90,7 @@ def test_multiline_docstring(self): ''' chunks = self.splitter._split_into_chunks(code) self.assertEqual(len(chunks), 2) - self.assertEqual(chunks[0].docstring.strip(), "Build initial structure\nWith multiple lines") + self.assertTrue("Build initial structure\nWith multiple lines" in chunks[0].code) def test_no_docstring(self): code = "x = 1\ny = 2" @@ -109,8 +124,8 @@ def test_empty_chunks(self): x = 1 ''' chunks = self.splitter._split_into_chunks(code) - self.assertEqual(len(chunks), 1) - self.assertEqual(chunks[0].docstring, "Task 2") + self.assertEqual(len(chunks), 2) + self.assertTrue("Task 2" in chunks[1].code) def test_code_with_strings(self): code = ''' diff --git a/src/tests/mcts/test_save_load_python_namespace.py b/src/tests/mcts/test_save_load_python_namespace.py new file mode 100644 index 000000000..a04f40103 --- /dev/null +++ b/src/tests/mcts/test_save_load_python_namespace.py @@ -0,0 +1,38 @@ +import unittest +import ast +from dataclasses import dataclass +from typing import List, Optional + +from datasetgen.mcts.chunked_mcts import ChunkedMCTS +from datasetgen.mcts.game_state import GameState +from factorio_instance import FactorioInstance + + +class TestSaveLoadPythonNamespace(unittest.TestCase): + """ + FactorioInstance exposes a Python namespace for the agent to persist variables. + These tests verify that the namespace can be saved and loaded correctly into new instances. + """ + def setUp(self): + self.instance = FactorioInstance(address='localhost', + bounding_box=200, + tcp_port=27015, + fast=True, + inventory={}) + + def test_save_load_namespace(self): + self.instance.eval('x=2') + game_state = GameState.from_instance(self.instance) + + self.instance = FactorioInstance(address='localhost', + bounding_box=200, + tcp_port=27015, + fast=True, + inventory={}) + + self.instance.reset(game_state) + self.instance.eval('assert x == 2') + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file From 77c21971f73f4fc9a30edbb2435cf966deb5af8c Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Tue, 19 Nov 2024 13:35:29 +0000 Subject: [PATCH 05/15] Add inventory and entities to the observation state if missing via modifying agent generated code. --- src/datasetgen/mcts/__main__.py | 8 +- src/datasetgen/mcts/chunked_mcts.py | 152 ++++++++++-------- src/datasetgen/mcts/conversation_formatter.py | 5 +- src/datasetgen/mcts/factorio_evaluator.py | 16 +- src/datasetgen/mcts/logger.py | 2 +- src/datasetgen/mcts/mcts.py | 6 +- src/datasetgen/mcts/program.py | 3 + src/tests/mcts/test_mcts_chunker.py | 85 ++++++++++ 8 files changed, 202 insertions(+), 75 deletions(-) diff --git a/src/datasetgen/mcts/__main__.py b/src/datasetgen/mcts/__main__.py index efbf58185..366a7fec4 100644 --- a/src/datasetgen/mcts/__main__.py +++ b/src/datasetgen/mcts/__main__.py @@ -61,7 +61,7 @@ async def main(): instance.speed(10) # Set the game speed to 10x normal speed for faster testing # Initialize FactorioEvaluator with the list of instances - evaluator = FactorioEvaluator(db_client, instances) + evaluator = FactorioEvaluator(db_client, instances, value_accrual_time=3) initial_state = GameState.from_instance(instances[0]) # Get execution directory from __file__ or other source @@ -77,15 +77,15 @@ async def main(): evaluator, system_prompt, initial_state, - version=4, - version_description="Step-wise evaluation / Execution results exclude entities and inventory", + version=5, + version_description="Step-wise evaluation / Errors not saved / Execution results exclude entities and inventory", formatter=StructurePreservingFormatter(planning=True)) print("Starting MCTS search...") best_programs = await mcts.search( n_iterations=500, samples_per_iteration=len(instances)-1, # One for each instance, minus a holdout. - skip_failures=False, + skip_failures=True, ) print("\nBest programs found:") diff --git a/src/datasetgen/mcts/chunked_mcts.py b/src/datasetgen/mcts/chunked_mcts.py index 3c9119434..7dbc30988 100644 --- a/src/datasetgen/mcts/chunked_mcts.py +++ b/src/datasetgen/mcts/chunked_mcts.py @@ -13,62 +13,60 @@ class ChunkedMCTS(MCTS): def _split_into_chunks(self, program_code: str) -> List[Program]: - chunks = [] + """Split the program code into chunks based on docstrings.""" + + program_code = program_code.replace("from factorio_instance import *", "").strip() lines = program_code.splitlines() - current_docstring = "" - current_lines = [] - current_start = 0 + chunks = [] module = ast.parse(program_code) - # Track docstring positions to identify chunk boundaries + # Find all docstrings and their positions docstring_positions = [] - for i, node in enumerate(module.body): + for node in module.body: if isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant): docstring_positions.append((node.lineno - 1, node.end_lineno - 1, node.value.s)) if not docstring_positions: return [] - # Process chunks between docstrings - for i, (start_pos, end_pos, docstring) in enumerate(docstring_positions): - # If this isn't the first chunk, append the previous one - if current_docstring and current_lines: + # Handle everything before first docstring + first_docstring_start = docstring_positions[0][0] + if first_docstring_start > 0: + preamble = lines[:first_docstring_start] + if any(line.strip() for line in preamble): # If there's any non-empty content chunks.append(Program( - code="\n".join(current_lines), + code='\n'.join(preamble), conversation=Conversation(messages=[]) )) - # Start new chunk - current_docstring = docstring - current_lines = [] + # Process each chunk + for i, (start_pos, end_pos, docstring) in enumerate(docstring_positions): + chunk_lines = [] - # Add any lines before the docstring that follow the previous chunk - if i > 0: - prev_end = docstring_positions[i - 1][1] - current_lines.extend(lines[prev_end + 1:start_pos]) + # Add docstring lines + chunk_lines.extend(lines[start_pos:end_pos + 1]) - # Add the docstring lines - current_lines.extend(lines[start_pos:end_pos + 1]) + # Add code lines until next docstring or end + if i < len(docstring_positions) - 1: + next_start = docstring_positions[i + 1][0] + chunk_lines.extend(lines[end_pos + 1:next_start]) + else: + # For last chunk, add all remaining lines + chunk_lines.extend(lines[end_pos + 1:]) - # If this is the last docstring, add all remaining lines - if i == len(docstring_positions) - 1: - current_lines.extend(lines[end_pos + 1:]) + # Create program for this chunk + if chunk_lines: chunks.append(Program( - code="\n".join(current_lines), + code='\n'.join(chunk_lines), conversation=Conversation(messages=[]) )) - else: - # Add lines until the next docstring - next_start = docstring_positions[i + 1][0] - current_lines.extend(lines[end_pos + 1:next_start]) return chunks async def _evaluate_chunks(self, chunks: List[Program], start_state: GameState, instance_id: int) \ -> Tuple[List[Program], float]: current_state = start_state - holdout_start = GameState.from_instance(self.evaluator.holdout) - self.evaluator.holdout.reset(holdout_start) + self.evaluator.holdout.reset(current_state) holdout_future = asyncio.create_task(self.evaluator._run_holdout()) entity_list = [] try: @@ -108,44 +106,72 @@ async def search(self, n_iterations: int, samples_per_iteration: int, skip_failu self.evaluator.set_sampling_status() raw_programs = await self._generate_programs_batch(conversation, samples_per_iteration) + # Process programs in parallel + eval_futures = [] for i, (program, chunks) in enumerate(raw_programs): instance_id = i % (len(self.evaluator.instances) - 1) + self.evaluator.instances[instance_id].reset(start_state) + self.evaluator.logger.update_instance(i, program_id=program.id, status="resetting") + + # Create evaluation future for this program's chunks + eval_futures.append(self._process_program_chunks( + program=program, + chunks=chunks, + start_state=start_state, + instance_id=instance_id, + parent_id=parent.id if parent else None, + skip_failures=skip_failures + )) + # Wait for all programs to complete + await asyncio.gather(*eval_futures) + self.evaluator.logger.update_progress() - self.evaluator.instances[instance_id].reset(start_state) # Reset to the start state before evaluating each chunk. - self.logger.update_instance(i, program_id=program.id, status="resetting") - - try: - evaluated_chunks, holdout, entity_list = await self._evaluate_chunks(chunks, start_state, instance_id) - last_chunk_id = parent.id if parent else None - last_conversation_stage = program.conversation - for chunk, entities in zip(evaluated_chunks, entity_list): - chunk_program = Program( - code=chunk.code, - conversation=program.conversation, - value=chunk.value - (holdout / len(evaluated_chunks)), # Distribute holdout value - state=chunk.state, - response=chunk.response, - version=self.version, - version_description=self.version_description, - parent_id=last_chunk_id, - token_usage=program.token_usage // len(evaluated_chunks), - completion_token_usage=program.completion_token_usage // len(evaluated_chunks), - prompt_token_usage=program.prompt_token_usage // len(evaluated_chunks) - ) - last_conversation_stage.add_result(chunk.code, chunk.value - (holdout / len(evaluated_chunks)), chunk.response, chunk.state, entities) - chunk_program.id = hash((chunk.code, json.dumps(chunk_program.conversation.model_dump()['messages']))) - chunk_program.conversation = last_conversation_stage - - last_chunk_id = chunk_program.id - if not skip_failures or chunk_program.value is not None: - await self.db.create_program(chunk_program) - - except Exception as e: - print(f"Failed to evaluate program: {str(e)}") - continue + async def _process_program_chunks(self, program: Program, chunks: List[Program], + start_state: GameState, instance_id: int, + parent_id: Optional[int], skip_failures: bool): + try: + evaluated_chunks, holdout, entity_list = await self._evaluate_chunks(chunks, start_state, instance_id) + last_chunk_id = parent_id + last_conversation_stage = program.conversation + + for chunk, entities in zip(evaluated_chunks, entity_list): + chunk_program = Program( + code=chunk.code, + conversation=program.conversation, + value=chunk.value - (holdout / len(evaluated_chunks)), + raw_reward=chunk.value, + holdout_value=holdout / len(evaluated_chunks), + state=chunk.state, + response=chunk.response, + version=self.version, + version_description=self.version_description, + parent_id=last_chunk_id, + token_usage=program.token_usage // len(evaluated_chunks), + completion_token_usage=program.completion_token_usage // len(evaluated_chunks), + prompt_token_usage=program.prompt_token_usage // len(evaluated_chunks) + ) + + last_conversation_stage.add_result( + chunk.code, + chunk.value - (holdout / len(evaluated_chunks)), + chunk.response, + chunk.state, + entities + ) - self.evaluator.logger.update_progress() + chunk_program.id = hash( + (chunk.code, json.dumps(chunk_program.conversation.model_dump()['messages']))) + chunk_program.conversation = last_conversation_stage + + if skip_failures and "error" in chunk.response.lower(): + break + + created = await self.db.create_program(chunk_program) + last_chunk_id = created.id + + except Exception as e: + print(f"Failed to evaluate program on instance {instance_id}: {str(e)}") async def _generate_programs_batch(self, conversation: Conversation, n_samples: int) -> List[Tuple[Program, List[Program]]]: programs = await super()._generate_programs_batch(conversation, n_samples) diff --git a/src/datasetgen/mcts/conversation_formatter.py b/src/datasetgen/mcts/conversation_formatter.py index 039169cb1..dd83e4ce1 100644 --- a/src/datasetgen/mcts/conversation_formatter.py +++ b/src/datasetgen/mcts/conversation_formatter.py @@ -7,8 +7,9 @@ PLANNING_ADDITION_PROMPT = \ """ -Identify an appropriate next task given the recent history and environment feedback. -Develop thorough step-by-step plan how you can achieve this task and then create the python script to achieve the task. +Your goal is to automate an increasingly complex factory process. +Identify an appropriate next task given the recent history and environment feedback to achieve this goal. +Develop thorough step-by-step plans for how you can achieve the next task and then create the python script to achieve the task. For your plan, follow this structure: 1) What entities are needed for the task 2) What entities do we have on the map, in different entity inventories or in our inventory diff --git a/src/datasetgen/mcts/factorio_evaluator.py b/src/datasetgen/mcts/factorio_evaluator.py index 00af76e87..3ef010aef 100644 --- a/src/datasetgen/mcts/factorio_evaluator.py +++ b/src/datasetgen/mcts/factorio_evaluator.py @@ -92,8 +92,20 @@ async def _evaluate_single(self, instance_id: int, program: Program, instance: F entities = instance.get_entities() final_inventory = instance.inspect_inventory() + # Check to see if the inventories are different + # If so, we put a hint in the code and result + if start_inventory.__dict__ != final_inventory.__dict__ and 'error' not in result.lower(): + program.code += '\nprint(f"Inventory changed to {inspect_inventory()}")' + result += f'\n('+str(len(program.code.split('\n')))+f': Inventory changed to {final_inventory},)' + + # Check to see if the entities are different + # If so, we put a hint in the code and result + if start_entities != entities and 'error' not in result.lower(): + program.code += '\nprint(f"Entities on the map: {get_entities()}")\n' + result += "\n("+str(len(program.code.split('\n')))+f': Entities on the map: {entities},)' + self.logger.update_instance(instance_id, status="accruing value") - await asyncio.sleep(10) + await asyncio.sleep(self.value_accrual_time) score, _ = instance.score() final_reward = score - initial_value @@ -134,7 +146,7 @@ async def _run_holdout(self) -> float: initial_value, _ = self.holdout.score() self.logger.update_instance(len(self.instances), status="accruing value") - await asyncio.sleep(10) + await asyncio.sleep(self.value_accrual_time) entities = self.holdout.get_entities() reward, _ = self.holdout.score() final_inventory = self.holdout.inspect_inventory() diff --git a/src/datasetgen/mcts/logger.py b/src/datasetgen/mcts/logger.py index 68ae6c401..e168488dc 100644 --- a/src/datasetgen/mcts/logger.py +++ b/src/datasetgen/mcts/logger.py @@ -104,7 +104,7 @@ def _generate_instance_panel(self, instance: InstanceMetrics) -> Panel: table.add_row("Total Programs:", str(instance.total_programs)) table.add_row("Last Update:", instance.last_update.strftime("%H:%M:%S")) table.add_row("Avg Time / Program:", Text(avg_time)) - table.add_row("# Entities:", Text(f"{instance.start_entities} -> {instance.final_entities}", style="cyan" if instance.start_entities >= instance.final_entities else "red")) + table.add_row("# Entities:", Text(f"{instance.start_entities} -> {instance.final_entities}", style="red" if instance.start_entities != instance.final_entities else "cyan")) table.add_row("# Inventory:", Text(f"{instance.start_inventory_count} -> {instance.final_inventory_count}")) return Panel( diff --git a/src/datasetgen/mcts/mcts.py b/src/datasetgen/mcts/mcts.py index 98d75b507..1b90718a2 100644 --- a/src/datasetgen/mcts/mcts.py +++ b/src/datasetgen/mcts/mcts.py @@ -59,15 +59,12 @@ def _extract_code_from_choice(self, choice) -> Optional[str]: code = self._verify_response_is_python(code) return code except Exception as e1: - print(f"Failed to extract code from choice: {str(e1)}") # Sometimes it samples a leading line, before writing unblocked python code. content = "\n".join(choice.message.content.split("\n")[1:]) try: code = self._verify_response_is_python(content) return code except Exception as e2: - print(f"Failed to extract code from choice after removing leading line: {str(e2)}") - try: code = content.split('from factorio_instance import *')[1].strip() code = self._verify_response_is_python(code) @@ -75,6 +72,8 @@ def _extract_code_from_choice(self, choice) -> Optional[str]: except Exception as e2: print(f"Failed to extract code from choice after removing leading line and factorio_instance import: {str(e2)}") return None + print(f"Failed to extract code from choice after removing leading line: {str(e2)}") + print(f"Failed to extract code from choice: {str(e1)}") @retry(wait=wait_exponential(multiplier=1, min=4, max=10)) async def _generate_programs_batch(self, conversation: Conversation, n_samples: int) -> List[Program]: @@ -89,6 +88,7 @@ async def _generate_programs_batch(self, conversation: Conversation, n_samples: # Single API call to generate n_samples completions response = self.llm.call( messages=formatted_messages, + max_tokens=2048, n=n_samples, temperature=0.7 # Adjust as needed ) diff --git a/src/datasetgen/mcts/program.py b/src/datasetgen/mcts/program.py index 7bfeb7b64..5242fe1f1 100644 --- a/src/datasetgen/mcts/program.py +++ b/src/datasetgen/mcts/program.py @@ -27,6 +27,9 @@ class Program(BaseModel): version: int = 1 version_description: str = "" + def __repr__(self): + return self.code + def get_uct(self, parent_visits: int, exploration_constant: float = 1.41) -> float: if self.visits == 0: return float('inf') diff --git a/src/tests/mcts/test_mcts_chunker.py b/src/tests/mcts/test_mcts_chunker.py index 8b726a536..1b1e4610e 100644 --- a/src/tests/mcts/test_mcts_chunker.py +++ b/src/tests/mcts/test_mcts_chunker.py @@ -13,6 +13,87 @@ class ProgramChunk: state: Optional['GameState'] = None reward: float = 0.0 +FULL_PROGRAM = \ +''' +from factorio_instance import * + +""" +Objective: Craft 8 automation science packs + +Planning: +1. Print recipe for automation science packs +2. Analyze current inventory +3. Craft iron gear wheels (we need 8) +4. Craft automation science packs (we need 8) +5. Verify the crafting process +""" + +""" +Step 1: Print recipe for automation science packs +""" +recipe = get_prototype_recipe(Prototype.AutomationSciencePack) +print("Automation Science Pack Recipe:") +print(f"Ingredients: {recipe.ingredients}") + +""" +Step 2: Analyze current inventory +""" +current_inventory = inspect_inventory() +print("Current inventory:") +for item, quantity in current_inventory.items(): + print(f"{item}: {quantity}") + +""" +Step 3: Craft iron gear wheels +""" +craft_item(Prototype.IronGearWheel, quantity=8) +print("Crafted 8 Iron Gear Wheels") + +""" +Step 4: Craft automation science packs +""" +craft_item(Prototype.AutomationSciencePack, quantity=8) +print("Crafted 8 Automation Science Packs") + +""" +Step 5: Verify the crafting process +""" +final_inventory = inspect_inventory() +automation_science_packs = final_inventory.get(Prototype.AutomationSciencePack, 0) +print(f"Final count of Automation Science Packs: {automation_science_packs}") + +assert automation_science_packs == 8, f"Failed to craft 8 Automation Science Packs. Current count: {automation_science_packs}" +print("Successfully crafted 8 Automation Science Packs!") +''' + +FULL_PROGRAM2= \ +''' +""" +Step 4: Extract steel plates +- Extract the steel plates from the second furnace +- Verify that we have at least 10 steel plates in our inventory +""" +# Extract steel plates from the second furnace +second_furnace = get_entity(Prototype.StoneFurnace, Position(x=-8.0, y=0.0)) +move_to(second_furnace.position) + +# Attempt to extract steel plates multiple times to ensure all are gathered +max_attempts = 5 +for _ in range(max_attempts): + extract_item(Prototype.SteelPlate, second_furnace.position, quantity=50) + steel_plates_in_inventory = inspect_inventory().get(Prototype.SteelPlate, 0) + if steel_plates_in_inventory >= 10: + break + sleep(5) # Wait a bit more if needed + +print(f"Extracted steel plates. Current inventory count: {steel_plates_in_inventory}") + +# Verify that we have at least 10 steel plates +assert steel_plates_in_inventory >= 10, f"Failed to extract required number of steel plates. Only {steel_plates_in_inventory} found in inventory." +print("Successfully extracted required number of steel plates") +print(f"Final steel plate count in inventory: {steel_plates_in_inventory}") +print("Steel plate extraction and verification completed successfully.") +''' class TestProgramChunkSplitter(unittest.TestCase): def setUp(self): @@ -139,6 +220,10 @@ def test_code_with_strings(self): self.assertTrue('x = "This is a string"' in chunks[0].code) self.assertTrue('y = """This is a multiline' in chunks[0].code) + def test_full_code(self): + chunks = self.splitter._split_into_chunks(FULL_PROGRAM2) + + pass if __name__ == '__main__': unittest.main() \ No newline at end of file From da26964bd5e7f1ec986b50a83b5d46b073f643b4 Mon Sep 17 00:00:00 2001 From: Mart Date: Tue, 19 Nov 2024 17:18:27 +0000 Subject: [PATCH 06/15] added the plan sampler --- .../auto_curriculum/plan_sampler.py | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 src/datasetgen/auto_curriculum/plan_sampler.py diff --git a/src/datasetgen/auto_curriculum/plan_sampler.py b/src/datasetgen/auto_curriculum/plan_sampler.py new file mode 100644 index 000000000..154c4a058 --- /dev/null +++ b/src/datasetgen/auto_curriculum/plan_sampler.py @@ -0,0 +1,55 @@ +from typing import Any +from llm_factory import LLMFactory + +class PlanSampler(): + def __init__(self, model, system_prompt_path, starting_scenarios_folder,): + + self.model = model + self.system_prompt_path = system_prompt_path + self.llm_factory = LLMFactory(model) + self.starting_scenarios_folder = starting_scenarios_folder + + + def get_unsupervised_objective(self, instance): + """ + Runs the traces, adds the plan to the prompt and saves the traces to the save path + """ + starting_inventory = instance.inspect_inventory() + self.init_system_prompt(instance) + mining_setup = self.get_mining_setup(instance) + messages = [{"role": "system", "content": self.system_prompt}] + user_message = f"Your starting inventory is {starting_inventory}. Your initial mining setup is: {mining_setup}. Create a useful task that you can carry out in the current game and the python script to achieve the task" + messages.append({"role": "user", "content": user_message}) + response = self.llm_factory.call(messages=messages, + model = self.model, + temperature=0.7, + max_tokens=4096, + stop_sequences = ["\n"]) + + + full_output = response.choices[0].message.content + full_output = full_output.lower().replace("sure! the task i will carry out is", "").strip() + if "." in full_output: + full_output = full_output.split(".")[0] + return full_output + + def get_mining_setup(self, instance): + mining_setup = instance.get_entities() + if len(mining_setup) == 0: + mining_setup = "There are no entities on the map" + else: + mining_setup = f"The following entities are on the map and can be used: {mining_setup}" + return mining_setup + + def init_system_prompt(self, instance): + api_description = instance.get_system_prompt() + system_prompt_path = self.system_prompt_path + # read in the system prompt + with open(system_prompt_path, "r") as f: + system_prompt = f.read() + self.system_prompt = system_prompt.format(schema=api_description) + + + def __call__(self, instance, game_state_str) -> Any: + instance.reset(game_state_str) + return self.get_unsupervised_objective(instance) \ No newline at end of file From da2b7f88ef1c3bfe07317dddc54889728bfecbd0 Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Tue, 19 Nov 2024 22:46:09 +0000 Subject: [PATCH 07/15] Ensure holdout runs for chunked mcts. Ensure all 3 servers process tasks.' --- src/datasetgen/mcts/__main__.py | 9 +++ src/datasetgen/mcts/chunked_mcts.py | 8 +- src/datasetgen/mcts/conversation_formatter.py | 12 +-- src/datasetgen/mcts/factorio_evaluator.py | 4 +- src/datasetgen/mcts/mcts.py | 7 +- src/datasetgen/mcts/plan_generator_1.py | 77 +++++++++++++++++++ src/llm_factory.py | 1 + src/tests/mcts/test_conversation_formatter.py | 10 +-- 8 files changed, 111 insertions(+), 17 deletions(-) create mode 100644 src/datasetgen/mcts/plan_generator_1.py diff --git a/src/datasetgen/mcts/__main__.py b/src/datasetgen/mcts/__main__.py index 366a7fec4..cdcab1af7 100644 --- a/src/datasetgen/mcts/__main__.py +++ b/src/datasetgen/mcts/__main__.py @@ -72,11 +72,20 @@ async def main(): print("Initializing MCTS...") + # Logit bias + logit_bias = { + "15714": -100, # 'LINE' + "145968": -100, # ' CUT' + "27": -100, # '<' + "20225": -100 #'/>' + } + mcts = ChunkedMCTS(llm, db_client, evaluator, system_prompt, initial_state, + logit_bias=logit_bias, version=5, version_description="Step-wise evaluation / Errors not saved / Execution results exclude entities and inventory", formatter=StructurePreservingFormatter(planning=True)) diff --git a/src/datasetgen/mcts/chunked_mcts.py b/src/datasetgen/mcts/chunked_mcts.py index 7dbc30988..894cf13ba 100644 --- a/src/datasetgen/mcts/chunked_mcts.py +++ b/src/datasetgen/mcts/chunked_mcts.py @@ -12,6 +12,11 @@ class ChunkedMCTS(MCTS): + + def __init__(self, *args, logit_bias: Optional[float] = None, **kwargs): + super().__init__(*args, **kwargs) + self.logit_bias = logit_bias + def _split_into_chunks(self, program_code: str) -> List[Program]: """Split the program code into chunks based on docstrings.""" @@ -174,7 +179,8 @@ async def _process_program_chunks(self, program: Program, chunks: List[Program], print(f"Failed to evaluate program on instance {instance_id}: {str(e)}") async def _generate_programs_batch(self, conversation: Conversation, n_samples: int) -> List[Tuple[Program, List[Program]]]: - programs = await super()._generate_programs_batch(conversation, n_samples) + # We generate one extra program in case there is an error in parsing one. This way we can always return n_samples programs to keep the servers occupied. + programs = (await super()._generate_programs_batch(conversation, n_samples+1, logit_bias=self.logit_bias))[:n_samples] chunked_programs = [] for i, program in enumerate(programs): diff --git a/src/datasetgen/mcts/conversation_formatter.py b/src/datasetgen/mcts/conversation_formatter.py index dd83e4ce1..db8d5846a 100644 --- a/src/datasetgen/mcts/conversation_formatter.py +++ b/src/datasetgen/mcts/conversation_formatter.py @@ -104,9 +104,9 @@ def summarize_code_block(code: str, start_line: int = None, end_line: int = None if not in_docstring: # Start of docstring if code_start is not None: if code_lines == 1: - result.append(f"") + result.append(f"") else: - result.append(f"") + result.append(f"") code_start = None code_lines = 0 in_docstring = True @@ -123,9 +123,9 @@ def summarize_code_block(code: str, start_line: int = None, end_line: int = None if stripped.startswith('#'): if code_start is not None: if code_lines == 1: - result.append(f"") + result.append(f"") else: - result.append(f"") + result.append(f"") code_start = None code_lines = 0 result.append(line) @@ -138,9 +138,9 @@ def summarize_code_block(code: str, start_line: int = None, end_line: int = None # Handle any remaining code section if code_start is not None: if code_lines == 1: - result.append(f"") + result.append(f"") else: - result.append(f"") + result.append(f"") return '\n'.join(result) diff --git a/src/datasetgen/mcts/factorio_evaluator.py b/src/datasetgen/mcts/factorio_evaluator.py index 3ef010aef..ca5c136ed 100644 --- a/src/datasetgen/mcts/factorio_evaluator.py +++ b/src/datasetgen/mcts/factorio_evaluator.py @@ -96,13 +96,13 @@ async def _evaluate_single(self, instance_id: int, program: Program, instance: F # If so, we put a hint in the code and result if start_inventory.__dict__ != final_inventory.__dict__ and 'error' not in result.lower(): program.code += '\nprint(f"Inventory changed to {inspect_inventory()}")' - result += f'\n('+str(len(program.code.split('\n')))+f': Inventory changed to {final_inventory},)' + result += f'\n'+str(len(program.code.split('\n')))+f': (Inventory changed to {final_inventory},)' # Check to see if the entities are different # If so, we put a hint in the code and result if start_entities != entities and 'error' not in result.lower(): program.code += '\nprint(f"Entities on the map: {get_entities()}")\n' - result += "\n("+str(len(program.code.split('\n')))+f': Entities on the map: {entities},)' + result += "\n"+str(len(program.code.split('\n')))+f': (Entities on the map: {entities},)' self.logger.update_instance(instance_id, status="accruing value") await asyncio.sleep(self.value_accrual_time) diff --git a/src/datasetgen/mcts/mcts.py b/src/datasetgen/mcts/mcts.py index 1b90718a2..869928c69 100644 --- a/src/datasetgen/mcts/mcts.py +++ b/src/datasetgen/mcts/mcts.py @@ -1,6 +1,6 @@ import asyncio import json -from typing import Optional, List +from typing import Optional, List, Dict from tenacity import wait_exponential, retry @@ -76,7 +76,7 @@ def _extract_code_from_choice(self, choice) -> Optional[str]: print(f"Failed to extract code from choice: {str(e1)}") @retry(wait=wait_exponential(multiplier=1, min=4, max=10)) - async def _generate_programs_batch(self, conversation: Conversation, n_samples: int) -> List[Program]: + async def _generate_programs_batch(self, conversation: Conversation, n_samples: int, logit_bias: Optional[Dict[str, float]] = None) -> List[Program]: """Generate multiple programs in a single API call using 'n' parameter""" formatted_messages = self.formatter.to_llm_messages( self.formatter.format_conversation(conversation) @@ -90,7 +90,8 @@ async def _generate_programs_batch(self, conversation: Conversation, n_samples: messages=formatted_messages, max_tokens=2048, n=n_samples, - temperature=0.7 # Adjust as needed + temperature=0.7, # Adjust as needed + logit_bias=logit_bias ) programs = [] diff --git a/src/datasetgen/mcts/plan_generator_1.py b/src/datasetgen/mcts/plan_generator_1.py new file mode 100644 index 000000000..34002e783 --- /dev/null +++ b/src/datasetgen/mcts/plan_generator_1.py @@ -0,0 +1,77 @@ +import asyncio +from typing import List, Dict + +from datasetgen.mcts.conversation import Conversation, Message +from datasetgen.mcts.program import Program +from llm_factory import LLMFactory + + +class PlanGenerator: + def __init__(self, llm_factory: 'LLMFactory'): + self.llm = llm_factory + + def _get_example_prompt(self) -> List[Dict[str, str]]: + return [ + {"role": "system", + "content": "You generate Factorio automation plans. Follow the format in the examples."}, + {"role": "user", "content": "Generate a plan"}, + {"role": "assistant", "content": '''""" +Objective: Craft 2 burner inserters +The final success should be checked by looking if the burner inserters are in inventory +""" + +""" +Planning: +1) Print recipes for the entities we need to craft +2) Gather resources (iron plates, stone) for crafting +3) Craft a burner mining drill +4) Place the burner mining drill on an iron ore patch +5) Place the stone furnace below the burner mining drill +6) Connect the burner mining drill to the stone furnace using a burner inserter +7) Fuel both the burner mining drill and stone furnace with coal +8) Verify the setup by checking if iron plates are being produced +"""'''}, + {"role": "user", "content": "Another plan"}, + {"role": "assistant", "content": '''""" +Objective: Craft and place a lab at the origin + +Planning: +1) Print recipe for Lab +2) Print recipes for intermediate products (iron gear wheels, transport belts, electronic circuits) +3) Print recipe for Assembling Machine 1 (optional) +4) Use existing Burner Mining Drill and Stone Furnace to produce iron plates for crafting +5) Craft intermediate products: iron gear wheels, transport belts, electronic circuits +6) Craft the Lab +7) Move to origin and place the Lab +"""'''}, + {"role": "user", "content": "Generate a plan for making transport belts"}, + {"role": "assistant", "content": '''""" +Objective: We need to craft 4 transport belts +Planning: We need to print out the recipe for transport belts and then craft them. We have the required resources in our inventory so we can directly craft the transport belts +"""'''} + ] + + async def generate_plan(self) -> str: + messages = self._get_example_prompt() + messages.append({"role": "user", "content": "Generate another automation plan"}) + + response = self.llm.call( + messages=messages, + temperature=0.7, + max_tokens=1024 + ) + + return response.choices[0].message.content + + async def generate_plans(self, n: int) -> List[str]: + plans = [] + for _ in range(n): + plan = await self.generate_plan() + print(plan) + plans.append(plan) + return plans + +if __name__ == "__main__": + llm = LLMFactory("ft:gpt-4o-2024-08-06:paperplane-ai:fact-self-gen-planning:AQzcPI91") + plan_generator = PlanGenerator(llm_factory=llm) + asyncio.run(plan_generator.generate_plans(n=3)) diff --git a/src/llm_factory.py b/src/llm_factory.py index ad557d836..fcb8bf470 100644 --- a/src/llm_factory.py +++ b/src/llm_factory.py @@ -110,6 +110,7 @@ def call(self, *args, **kwargs): max_tokens = kwargs.get('max_tokens', 2048), temperature=kwargs.get('temperature', 0.3), messages=kwargs.get('messages', None), + logit_bias=kwargs.get('logit_bias', None), n=n_samples, # Use requested number of samples #stop=["\n\n"],#, "\n#"], #presence_penalty=1, diff --git a/src/tests/mcts/test_conversation_formatter.py b/src/tests/mcts/test_conversation_formatter.py index 2aff35637..4dc6701e7 100644 --- a/src/tests/mcts/test_conversation_formatter.py +++ b/src/tests/mcts/test_conversation_formatter.py @@ -51,7 +51,7 @@ def test_code_summariser(self): summarized = CodeProcessor.summarize_code_block(code_block, preserve_comments=True) - self.assertEqual("# Gather iron ore\n\n# Construct stone furnace", summarized) + self.assertEqual("# Gather iron ore\n\n# Construct stone furnace", summarized) summarized2 = CodeProcessor.summarize_code_block( "# Gather iron ore\n# Gather more iron ore\nprint(0)\nprint(1)\n# Construct stone furnace", preserve_comments=True) @@ -73,7 +73,7 @@ def test_format_conversation(self): self.assertEqual(assistant1.role, "assistant") expected_summary = ( "# Gather iron ore\n" - "\n" + "\n" "# Construct stone furnace" ) self.assertEqual(assistant1.content, expected_summary) @@ -108,9 +108,9 @@ def test_format_single_message(self): formatted = self.formatter.format_message(message, is_last=False) expected = ( "# First task\n" - "\n" + "\n" "# Second task\n" - "" + "" ) self.assertEqual(formatted.content, expected) self.assertEqual(formatted.metadata, {"summarized": True}) @@ -132,7 +132,7 @@ def test_empty_code_blocks(self): formatted = self.formatter.format_message(message, is_last=False) self.assertEqual( formatted.content, - "" + "" ) def test_docstring_code_summariser(self): From 12461cb578e106d09091afb019889a641b99b5f7 Mon Sep 17 00:00:00 2001 From: Mart Date: Tue, 19 Nov 2024 23:13:55 +0000 Subject: [PATCH 08/15] added example how to use --- .../auto_curriculum/plan_sampler.py | 39 ++++++++++++++++++- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/src/datasetgen/auto_curriculum/plan_sampler.py b/src/datasetgen/auto_curriculum/plan_sampler.py index 154c4a058..8e6364c33 100644 --- a/src/datasetgen/auto_curriculum/plan_sampler.py +++ b/src/datasetgen/auto_curriculum/plan_sampler.py @@ -1,8 +1,12 @@ from typing import Any from llm_factory import LLMFactory +import os +from datasetgen.auto_curriculum.dataset_utils import instantiate_the_map, initialise_starting_scenario +from datasetgen.mcts.game_state import GameState +from factorio_instance import FactorioInstance class PlanSampler(): - def __init__(self, model, system_prompt_path, starting_scenarios_folder,): + def __init__(self, model, system_prompt_path, starting_scenarios_folder): self.model = model self.system_prompt_path = system_prompt_path @@ -52,4 +56,35 @@ def init_system_prompt(self, instance): def __call__(self, instance, game_state_str) -> Any: instance.reset(game_state_str) - return self.get_unsupervised_objective(instance) \ No newline at end of file + return self.get_unsupervised_objective(instance) + + def get_game_state(self, instance, starting_scenario_name): + # gets starting scenario details + starting_scenario_path = os.path.join(self.starting_scenarios_folder, starting_scenario_name) + starting_scenario = initialise_starting_scenario( + starting_scenario_path) # Gets the starting scenario details + # instantiate the map + result = instantiate_the_map(starting_scenario, instance, self.starting_scenarios_folder) + if not result["success"]: + print(f"Error in starting scenario: {result['error']}") + return False + # get the game state + game_state = GameState.from_instance(instance) + return game_state + +if __name__ == "__main__": + prompt_path = "prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md" + model_path = "path_to_model" + starting_scenario_folder = "skills/data_scenarios/starting_scenarios" + sampler = PlanSampler(model_path, prompt_path, starting_scenario_folder) + starting_scenario = "ft_random_chest_furnace_placement_with_mining_entities" + instance = FactorioInstance(address='localhost', + bounding_box=200, + tcp_port=27015, + fast=True, + #cache_scripts=False, + inventory={}) + + game_state = sampler.get_game_state(starting_scenario) + objective = sampler(instance, game_state) + From e8ed817f99691f37a80c565ea4fbbc40caec42c1 Mon Sep 17 00:00:00 2001 From: Mart Date: Tue, 19 Nov 2024 23:44:18 +0000 Subject: [PATCH 09/15] added correct stop sequence parsing --- src/llm_factory.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/llm_factory.py b/src/llm_factory.py index f1da2542a..803f5f2e3 100644 --- a/src/llm_factory.py +++ b/src/llm_factory.py @@ -104,12 +104,13 @@ def call(self, *args, **kwargs): stream=False) else: client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) + stop_sequences = kwargs.get('stop_sequences', None) assert "messages" in kwargs, "You must provide a list of messages to the model." return client.chat.completions.create(model = model_to_use, max_tokens = kwargs.get('max_tokens', 2048), temperature=kwargs.get('temperature', 0.3), messages=kwargs.get('messages', None), - #stop=["\n\n"],#, "\n#"], + stop=stop_sequences,#, "\n#"], #presence_penalty=1, #frequency_penalty=0.6, stream=False) From fc1faa83e4921690ddc1a12f445e57cd66490c95 Mon Sep 17 00:00:00 2001 From: Mart Date: Wed, 20 Nov 2024 10:03:30 +0000 Subject: [PATCH 10/15] added planning --- .../auto_curriculum/plan_sampler.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/datasetgen/auto_curriculum/plan_sampler.py b/src/datasetgen/auto_curriculum/plan_sampler.py index 8e6364c33..14d027c26 100644 --- a/src/datasetgen/auto_curriculum/plan_sampler.py +++ b/src/datasetgen/auto_curriculum/plan_sampler.py @@ -12,6 +12,14 @@ def __init__(self, model, system_prompt_path, starting_scenarios_folder): self.system_prompt_path = system_prompt_path self.llm_factory = LLMFactory(model) self.starting_scenarios_folder = starting_scenarios_folder + self.planning_addition_for_prompt = """ +First bring out a thorough step-by-step plan how you can achieve this task and then create the python script to achieve the task. +For your plan, follow this structure: +1) What entities are needed for the task +2) What entities do we have on the map, in different entity inventories or in our inventory +3) What entities are we missing for the task +4) Execution -- Taking into account 1,2 and 3, what steps do we need to take to successfully carry out the task +""" def get_unsupervised_objective(self, instance): @@ -23,18 +31,19 @@ def get_unsupervised_objective(self, instance): mining_setup = self.get_mining_setup(instance) messages = [{"role": "system", "content": self.system_prompt}] user_message = f"Your starting inventory is {starting_inventory}. Your initial mining setup is: {mining_setup}. Create a useful task that you can carry out in the current game and the python script to achieve the task" + user_message += f"\n{self.planning_addition_for_prompt}" messages.append({"role": "user", "content": user_message}) response = self.llm_factory.call(messages=messages, model = self.model, temperature=0.7, max_tokens=4096, - stop_sequences = ["\n"]) + stop_sequences = ["```"]) full_output = response.choices[0].message.content - full_output = full_output.lower().replace("sure! the task i will carry out is", "").strip() - if "." in full_output: - full_output = full_output.split(".")[0] + full_output = full_output.strip() + new_line_idx = full_output.rfind("\n") + full_output = full_output[:new_line_idx].replace("Sure!", "").strip() return full_output def get_mining_setup(self, instance): @@ -85,6 +94,6 @@ def get_game_state(self, instance, starting_scenario_name): #cache_scripts=False, inventory={}) - game_state = sampler.get_game_state(starting_scenario) + game_state = sampler.get_game_state(instance, starting_scenario) objective = sampler(instance, game_state) From 56557b2a5612b86482bac1d520cdb684c85e37a1 Mon Sep 17 00:00:00 2001 From: Mart Date: Wed, 20 Nov 2024 10:05:05 +0000 Subject: [PATCH 11/15] added docstring --- src/datasetgen/auto_curriculum/plan_sampler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/datasetgen/auto_curriculum/plan_sampler.py b/src/datasetgen/auto_curriculum/plan_sampler.py index 14d027c26..b49c75f52 100644 --- a/src/datasetgen/auto_curriculum/plan_sampler.py +++ b/src/datasetgen/auto_curriculum/plan_sampler.py @@ -44,6 +44,7 @@ def get_unsupervised_objective(self, instance): full_output = full_output.strip() new_line_idx = full_output.rfind("\n") full_output = full_output[:new_line_idx].replace("Sure!", "").strip() + full_output = f'"""\n{full_output}\n"""' return full_output def get_mining_setup(self, instance): From bd20b8171ba93af0a7fb38d114e6a9c049ad0545 Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Thu, 21 Nov 2024 17:18:15 +0000 Subject: [PATCH 12/15] Fixed error serialization bug. Added objective seeding. --- src/actions/connect_entities.lua | 2 +- src/actions/craft_item.lua | 18 ++-- src/actions/get_entity.lua | 2 +- src/actions/place_entity_next_to.lua | 2 +- src/actions/place_object.lua | 6 +- .../auto_curriculum/plan_sampler.py | 80 ++++++++++++---- src/datasetgen/mcts/__main__.py | 92 ++++++++++++++++++- src/datasetgen/mcts/chunked_mcts.py | 3 +- src/datasetgen/mcts/conversation_formatter.py | 4 +- src/datasetgen/mcts/db_client.py | 58 +++++++++++- src/datasetgen/mcts/factorio_evaluator.py | 6 +- src/datasetgen/mcts/schema.sql | 4 +- src/factorio_rcon_utils.py | 6 +- src/factorio_types.py | 2 +- src/llm_factory.py | 1 + .../system_message_policy_self_gen.md | 4 +- src/tests/actions/test_craft.py | 8 +- src/tests/complex/test_edge_cases.py | 2 - 18 files changed, 245 insertions(+), 55 deletions(-) diff --git a/src/actions/connect_entities.lua b/src/actions/connect_entities.lua index 86c86fa92..c2fd7e00a 100644 --- a/src/actions/connect_entities.lua +++ b/src/actions/connect_entities.lua @@ -485,7 +485,7 @@ global.actions.connect_entities = function(player_index, source_x, source_y, tar game.print("Required count: " .. required_count) game.print("Available count: " .. number_of_connection_entities) if number_of_connection_entities < required_count then - error("Player does not have enough " .. connection_type .. " in their inventory to complete this connection. Required number: " .. required_count .. ", Available in inventory: " .. number_of_connection_entities) + error("\"Player does not have enough " .. connection_type .. " in their inventory to complete this connection. Required number: " .. required_count .. ", Available in inventory: " .. number_of_connection_entities.."\"") end result = connect_entities(player_index, source_x, source_y, target_x, target_y, path_handle, connection_type, false) end diff --git a/src/actions/craft_item.lua b/src/actions/craft_item.lua index 571db6555..4cf5bc796 100644 --- a/src/actions/craft_item.lua +++ b/src/actions/craft_item.lua @@ -139,10 +139,10 @@ global.actions.craft_item = function(player_index, entity, count) if total_crafted >= count or (not global.fast and total_crafted > 0) then return count elseif total_crafted > 0 then - error("Successfully crafted " .. total_crafted .."x but failed to craft " - .. (count-total_crafted) .. "x " .. entity.." because ".. reason) + error("\"Successfully crafted " .. total_crafted .."x but failed to craft " + .. (count-total_crafted) .. "x " .. entity.." because ".. reason.."\"") else - error("Failed to craft " .. count .. "x " .. entity.." because "..reason) + error("\"Failed to craft " .. count .. "x " .. entity.." because "..reason.."\"") end end @@ -303,10 +303,10 @@ global.actions.craft_item3 = function(player_index, entity, count) if total_crafted >= count or (not global.fast and total_crafted > 0) then return count elseif total_crafted > 0 then - error("Successfully crafted " .. total_crafted .."x but failed to craft " - .. (count-total_crafted) .. "x " .. entity.." because ".. reason) + error("\"Successfully crafted " .. total_crafted .."x but failed to craft " + .. (count-total_crafted) .. "x " .. entity.." because ".. reason.."\"") else - error("Failed to craft " .. count .. "x " .. entity.." because "..reason) + error("\"Failed to craft " .. count .. "x " .. entity.." because "..reason.."\"") end end @@ -391,9 +391,9 @@ global.actions.craft_item2 = function(player_index, entity, count) game.print("Crafted x"..count.." "..entity) return count elseif total_crafted > 0 then - error("Successfully crafted " .. total_crafted .."x but failed to craft " - .. (count-total_crafted) .. "x " .. entity.." because ".. reason) + error("\"Successfully crafted " .. total_crafted .."x but failed to craft " + .. (count-total_crafted) .. "x " .. entity.." because ".. reason.."\"") else - error("Failed to craft " .. count .. "x_" .. entity.." because "..reason) + error("\"Failed to craft " .. count .. "x_" .. entity.." because "..reason.."\"") end end \ No newline at end of file diff --git a/src/actions/get_entity.lua b/src/actions/get_entity.lua index 4b9c44bf2..ac602fb10 100644 --- a/src/actions/get_entity.lua +++ b/src/actions/get_entity.lua @@ -29,6 +29,6 @@ global.actions.get_entity = function(player_index, entity, x, y) --local entity_json = game.table_to_json(serialized)-- game.table_to_json(entity return serialized else - error("No entity of type " .. entity .. " found at the specified position.") + error("\"No entity of type " .. entity .. " found at the specified position.\"") end end diff --git a/src/actions/place_entity_next_to.lua b/src/actions/place_entity_next_to.lua index 3942e2516..9d5cc7d3e 100644 --- a/src/actions/place_entity_next_to.lua +++ b/src/actions/place_entity_next_to.lua @@ -181,7 +181,7 @@ global.actions.place_entity_next_to = function(player_index, entity, ref_x, ref_ else colliding_entity_name = table.concat(colliding_entity_names, ", ", 1, #colliding_entity_names - 1) .. ", and " .. colliding_entity_names[#colliding_entity_names] end - error("A " .. colliding_entity_name .. " already exists at the new position " .. serpent.line(new_position) .. ".") + error("\"A " .. colliding_entity_name .. " already exists at the new position " .. serpent.line(new_position) .. ".\"") end end diff --git a/src/actions/place_object.lua b/src/actions/place_object.lua index d00e12eaf..7cb9f92c1 100644 --- a/src/actions/place_object.lua +++ b/src/actions/place_object.lua @@ -10,11 +10,11 @@ global.actions.place_entity = function(player_index, entity, direction, x, y, ex local distance = math.sqrt(dx * dx + dy * dy) if distance > max_distance then - error("The target position is too far away to place the entity. The player position is " .. + error("\"The target position is too far away to place the entity. The player position is " .. player.position.x .. ", " .. player.position.y .. " and the target position is " .. x .. ", " .. y .. ". The distance is " .. distance .. - " and the max distance is " .. max_distance .. ". Move closer.") + " and the max distance is " .. max_distance .. ". Move closer.\"") end end @@ -232,7 +232,7 @@ global.actions.place_entity2 = function(player_index, entity, direction, x, y, e local distance = math.sqrt(dx * dx + dy * dy) if distance > max_distance then - error("The target position is too far away to place the entity. The player position is " .. player.position.x .. ", " .. player.position.y .. " and the target position is " .. x .. ", " .. y .. ". The distance is " .. distance .. " and the max distance is " .. max_distance .. ". Move closer.") + error("\"The target position is too far away to place the entity. The player position is " .. player.position.x .. ", " .. player.position.y .. " and the target position is " .. x .. ", " .. y .. ". The distance is " .. distance .. " and the max distance is " .. max_distance .. ". Move closer.\"") end diff --git a/src/datasetgen/auto_curriculum/plan_sampler.py b/src/datasetgen/auto_curriculum/plan_sampler.py index b49c75f52..c271c5212 100644 --- a/src/datasetgen/auto_curriculum/plan_sampler.py +++ b/src/datasetgen/auto_curriculum/plan_sampler.py @@ -1,4 +1,8 @@ +from pyexpat.errors import messages from typing import Any + +from datasetgen.mcts.conversation import Conversation +from datasetgen.mcts.program import Program from llm_factory import LLMFactory import os from datasetgen.auto_curriculum.dataset_utils import instantiate_the_map, initialise_starting_scenario @@ -19,6 +23,32 @@ def __init__(self, model, system_prompt_path, starting_scenarios_folder): 2) What entities do we have on the map, in different entity inventories or in our inventory 3) What entities are we missing for the task 4) Execution -- Taking into account 1,2 and 3, what steps do we need to take to successfully carry out the task + +The format should look like this: + +''' +Objective: XXX + +Planning: +1) Entities needed for the task: +... + +2) Entities we have: +From the initial inventory: +... + +Entities on the map: +... + +3) Entities we are missing: +... + +4) Execution steps: +- ... +- ... +etc... +''' + """ @@ -36,7 +66,7 @@ def get_unsupervised_objective(self, instance): response = self.llm_factory.call(messages=messages, model = self.model, temperature=0.7, - max_tokens=4096, + max_tokens=256, stop_sequences = ["```"]) @@ -45,7 +75,7 @@ def get_unsupervised_objective(self, instance): new_line_idx = full_output.rfind("\n") full_output = full_output[:new_line_idx].replace("Sure!", "").strip() full_output = f'"""\n{full_output}\n"""' - return full_output + return full_output, response def get_mining_setup(self, instance): mining_setup = instance.get_entities() @@ -66,8 +96,12 @@ def init_system_prompt(self, instance): def __call__(self, instance, game_state_str) -> Any: instance.reset(game_state_str) - return self.get_unsupervised_objective(instance) - + objective, response = self.get_unsupervised_objective(instance) + try: + return objective.split("'''")[1].strip(), response + except: + return objective.split('"""')[2].strip(), response + def get_game_state(self, instance, starting_scenario_name): # gets starting scenario details starting_scenario_path = os.path.join(self.starting_scenarios_folder, starting_scenario_name) @@ -82,19 +116,27 @@ def get_game_state(self, instance, starting_scenario_name): game_state = GameState.from_instance(instance) return game_state -if __name__ == "__main__": - prompt_path = "prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md" - model_path = "path_to_model" - starting_scenario_folder = "skills/data_scenarios/starting_scenarios" - sampler = PlanSampler(model_path, prompt_path, starting_scenario_folder) - starting_scenario = "ft_random_chest_furnace_placement_with_mining_entities" - instance = FactorioInstance(address='localhost', - bounding_box=200, - tcp_port=27015, - fast=True, - #cache_scripts=False, - inventory={}) - - game_state = sampler.get_game_state(instance, starting_scenario) - objective = sampler(instance, game_state) + +# if __name__ == "__main__": +# prompt_path = "../../prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md" +# model_path = "ft:gpt-4o-2024-08-06:paperplane-ai:fact-self-gen-planning:AQzcPI91" +# starting_scenario_folder = "../../skills/data_scenarios/starting_scenarios" +# sampler = PlanSampler(model_path, prompt_path, starting_scenario_folder) +# starting_scenario = "ft_random_chest_furnace_placement_with_mining_entities" +# instance = FactorioInstance(address='localhost', +# bounding_box=200, +# tcp_port=27015, +# fast=True, +# #cache_scripts=False, +# inventory={}) +# +# game_state = sampler.get_game_state(instance, starting_scenario) +# objective = sampler(instance, game_state) +# +# program = Program(code=objective, +# conversation=Conversation(messages=[]), +# value=10, +# raw_reward=10, +# response="") +# pass \ No newline at end of file diff --git a/src/datasetgen/mcts/__main__.py b/src/datasetgen/mcts/__main__.py index cdcab1af7..313585b69 100644 --- a/src/datasetgen/mcts/__main__.py +++ b/src/datasetgen/mcts/__main__.py @@ -1,6 +1,11 @@ +import json import os +import statistics +from datasetgen.auto_curriculum.plan_sampler import PlanSampler from datasetgen.mcts.chunked_mcts import ChunkedMCTS +from datasetgen.mcts.conversation import Conversation, Message +from datasetgen.mcts.program import Program os.environ["FORCE_COLOR"] = "1" os.environ["TERM"] = "xterm-256color" @@ -13,7 +18,7 @@ from rich import print from rich.console import Console from cluster.local.cluster_ips import get_local_container_ips -from datasetgen.mcts.conversation_formatter import StructurePreservingFormatter +from datasetgen.mcts.conversation_formatter import StructurePreservingFormatter, PLANNING_ADDITION_PROMPT from datasetgen.mcts.db_client import DBClient from datasetgen.mcts.factorio_evaluator import FactorioEvaluator from datasetgen.mcts.game_state import GameState @@ -46,9 +51,76 @@ def create_parallel_instances() -> List[FactorioInstance]: return instances + +async def get_seed_programs( + mcts: ChunkedMCTS, + plan_sampler: 'PlanSampler', + n_seeds: int = 100, +) -> List[Program]: + existing_rewards = await mcts.db.get_all_program_rewards(version=mcts.version) + # filter out rewards < 0 + existing_rewards = list(filter(lambda x: x >= 0, existing_rewards)) + default_reward = (max(existing_rewards)-min(existing_rewards))/2 if existing_rewards else 10.0 + seeded_programs: List[Program] = [] + instance = mcts.evaluator.instances[0] + + # Get all available scenarios + scenarios = [f for f in os.listdir(plan_sampler.starting_scenarios_folder) + if os.path.isdir(os.path.join(plan_sampler.starting_scenarios_folder, f))] + + seeds_per_scenario = n_seeds // len(scenarios) + remaining_seeds = n_seeds % len(scenarios) + + for scenario in scenarios: + num_samples = seeds_per_scenario + (1 if remaining_seeds > 0 else 0) + remaining_seeds -= 1 + + for _ in range(num_samples): + game_state = plan_sampler.get_game_state(instance, scenario) + if not game_state: + continue + + objective, response = plan_sampler(instance, game_state) + + if len(objective) < 100: + continue + + + + conversation = Conversation(messages=[ + Message(role="system", content=mcts.system_prompt), + Message(role="user", + content=f"Inventory: {json.dumps(game_state.inventory.__dict__)}\n\n{PLANNING_ADDITION_PROMPT}"), + Message(role="assistant", content=objective) + ]) + messages = conversation.model_dump()['messages'] + if not objective.strip().startswith('"""'): + objective = '"""\n'+objective + + program = Program( + id=hash((objective, json.dumps(messages))), + code=objective, + conversation=conversation, + value=default_reward, + state=game_state, + version=mcts.version, + version_description=mcts.version_description, + token_usage=response.usage.total_tokens if hasattr(response, 'usage') else None, + completion_token_usage=response.usage.completion_tokens if hasattr(response, 'usage') else None, + prompt_token_usage=response.usage.prompt_tokens if hasattr(response, 'usage') else None, + ) + seeded_programs.append(program) + + return seeded_programs + async def main(): + model = "ft:gpt-4o-2024-08-06:paperplane-ai:fact-self-gen-planning:AQzcPI91" + prompt_path = "../../prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md" + version = 8 + version_description = "Seeded / No planning prompt in user messages / Step-wise evaluation / Errors not saved" + # Initialize components - llm = LLMFactory("ft:gpt-4o-2024-08-06:paperplane-ai:fact-self-gen-planning:AQzcPI91") + llm = LLMFactory(model) db_client = DBClient(host=os.getenv("SKILLS_DB_HOST"), port=os.getenv("SKILLS_DB_PORT"), dbname=os.getenv("SKILLS_DB_NAME"), @@ -80,21 +152,31 @@ async def main(): "20225": -100 #'/>' } + + mcts = ChunkedMCTS(llm, db_client, evaluator, system_prompt, initial_state, logit_bias=logit_bias, - version=5, - version_description="Step-wise evaluation / Errors not saved / Execution results exclude entities and inventory", + version=version, + version_description=version_description, formatter=StructurePreservingFormatter(planning=True)) + starting_scenario_folder = "../../skills/data_scenarios/starting_scenarios" + sampler = PlanSampler(model, prompt_path, starting_scenario_folder) + + print("Sampling seed scenarios...") + seeded_programs = await get_seed_programs(mcts, sampler, n_seeds=0) + for program in seeded_programs: + await db_client.create_program(program) + print("Starting MCTS search...") best_programs = await mcts.search( n_iterations=500, samples_per_iteration=len(instances)-1, # One for each instance, minus a holdout. - skip_failures=True, + skip_failures=False, ) print("\nBest programs found:") diff --git a/src/datasetgen/mcts/chunked_mcts.py b/src/datasetgen/mcts/chunked_mcts.py index 894cf13ba..cc82e97ff 100644 --- a/src/datasetgen/mcts/chunked_mcts.py +++ b/src/datasetgen/mcts/chunked_mcts.py @@ -1,7 +1,6 @@ import ast import asyncio import json -from dataclasses import dataclass from typing import List, Tuple, Optional from datasetgen.mcts.conversation import Conversation, Message @@ -114,7 +113,7 @@ async def search(self, n_iterations: int, samples_per_iteration: int, skip_failu # Process programs in parallel eval_futures = [] for i, (program, chunks) in enumerate(raw_programs): - instance_id = i % (len(self.evaluator.instances) - 1) + instance_id = i % (len(self.evaluator.instances)) self.evaluator.instances[instance_id].reset(start_state) self.evaluator.logger.update_instance(i, program_id=program.id, status="resetting") diff --git a/src/datasetgen/mcts/conversation_formatter.py b/src/datasetgen/mcts/conversation_formatter.py index db8d5846a..64cfe8e9b 100644 --- a/src/datasetgen/mcts/conversation_formatter.py +++ b/src/datasetgen/mcts/conversation_formatter.py @@ -207,8 +207,8 @@ def format_message(self, message: Message, is_last: bool = False) -> Optional[Me if "Execution result:" in content: result = content.split("Execution result:")[1].split("Updated state:")[0] content = f"Execution result:\n{result.strip()}" - if self.planning: - content = PLANNING_ADDITION_PROMPT + '\n' + content + #if self.planning: + # content = PLANNING_ADDITION_PROMPT + '\n' + content return Message( role="user", content=content diff --git a/src/datasetgen/mcts/db_client.py b/src/datasetgen/mcts/db_client.py index 7fdb1653e..292236b22 100644 --- a/src/datasetgen/mcts/db_client.py +++ b/src/datasetgen/mcts/db_client.py @@ -1,5 +1,5 @@ import json -from typing import Optional, Dict, Any +from typing import Optional, Dict, Any, List import psycopg2 import tenacity @@ -20,7 +20,7 @@ async def create_program(self, program: Program) -> Program: INSERT INTO programs (code, value, visits, parent_id, state_json, conversation_json, completion_token_usage, prompt_token_usage, token_usage, response, holdout_value, raw_reward, version, version_description) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING id, created_at - """, (program.code, program.value, program.visits, program.parent_id, + """, (program.code, program.value, 0, program.parent_id, program.state.to_raw() if program.state else None, json.dumps(program.conversation.dict()), program.completion_token_usage, @@ -42,6 +42,30 @@ async def create_program(self, program: Program) -> Program: print(e) raise e + + @tenacity.retry(retry=retry_if_exception_type((psycopg2.OperationalError, psycopg2.InterfaceError)), + wait=wait_exponential(multiplier=1, min=4, max=10)) + async def get_all_program_rewards(self, version: int = None) -> List[float]: + """Get all program rewards for a given version.""" + query = """ + SELECT value + FROM programs + WHERE value IS NOT NULL + """ + + if version is not None: + query += f" AND version = {version}" + + + try: + with self.conn.cursor() as cur: + cur.execute(query.strip()) + results = cur.fetchall() + return [row[0] for row in results] + except Exception as e: + print(f"Error fetching program rewards: {e}") + return [] + @tenacity.retry(retry=retry_if_exception_type((psycopg2.OperationalError, psycopg2.InterfaceError)), wait=wait_exponential(multiplier=1, min=4, max=10)) async def sample_parent(self, version=1) -> Optional[Program]: with self.conn.cursor() as cur: @@ -56,6 +80,36 @@ async def sample_parent(self, version=1) -> Optional[Program]: row = cur.fetchone() return Program.from_row(dict(zip([desc[0] for desc in cur.description], row))) + @tenacity.retry(retry=retry_if_exception_type((psycopg2.OperationalError, psycopg2.InterfaceError)), + wait=wait_exponential(multiplier=1, min=4, max=10)) + async def get_parent_visit_stats(self, version: int = None) -> Dict[str, float]: + """Get statistics about parent visit counts""" + query = """ + SELECT + AVG(visits) as avg_visits, + MIN(visits) as min_visits, + MAX(visits) as max_visits, + PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY visits) as median_visits + FROM programs + WHERE visits > 0 + """ + + if version is not None: + query += f" AND version = {version}" + + try: + with self.conn.cursor() as cur: + cur.execute(query) + result = cur.fetchone() + return { + 'avg_visits': result[0], + 'min_visits': result[1], + 'max_visits': result[2], + 'median_visits': result[3] + } + except Exception as e: + print(f"Error fetching visit statistics: {e}") + return {} async def update_program(self, program_id: int, updates: Dict[str, Any]) -> Program: with self.conn.cursor() as cur: diff --git a/src/datasetgen/mcts/factorio_evaluator.py b/src/datasetgen/mcts/factorio_evaluator.py index ca5c136ed..9cc2d604a 100644 --- a/src/datasetgen/mcts/factorio_evaluator.py +++ b/src/datasetgen/mcts/factorio_evaluator.py @@ -96,13 +96,13 @@ async def _evaluate_single(self, instance_id: int, program: Program, instance: F # If so, we put a hint in the code and result if start_inventory.__dict__ != final_inventory.__dict__ and 'error' not in result.lower(): program.code += '\nprint(f"Inventory changed to {inspect_inventory()}")' - result += f'\n'+str(len(program.code.split('\n')))+f': (Inventory changed to {final_inventory},)' + result += f'\n'+str(len(program.code.split('\n')))+f': (\'Inventory changed to {final_inventory}\',)' # Check to see if the entities are different # If so, we put a hint in the code and result if start_entities != entities and 'error' not in result.lower(): program.code += '\nprint(f"Entities on the map: {get_entities()}")\n' - result += "\n"+str(len(program.code.split('\n')))+f': (Entities on the map: {entities},)' + result += "\n"+str(len(program.code.split('\n')))+f': (\'Entities on the map: {entities}\',)' self.logger.update_instance(instance_id, status="accruing value") await asyncio.sleep(self.value_accrual_time) @@ -117,6 +117,7 @@ async def _evaluate_single(self, instance_id: int, program: Program, instance: F raw_reward=final_reward, final_entities=len(entities), start_entities=len(start_entities), + total_programs=self.logger.instances[instance_id].total_programs + 1, start_inventory_count=sum([v for k, v in start_inventory.__dict__.items() if v > 0]), final_inventory_count=sum([v for k, v in final_inventory.__dict__.items() if v > 0]) ) @@ -155,6 +156,7 @@ async def _run_holdout(self) -> float: status="accrued value", final_entities=len(entities), start_entities=len(initial_entities), + total_programs=self.logger.instances[len(self.instances)].total_programs + 1, start_inventory_count=sum([v for k, v in start_inventory.__dict__.items() if v > 0]), final_inventory_count=sum([v for k, v in final_inventory.__dict__.items() if v > 0])) diff --git a/src/datasetgen/mcts/schema.sql b/src/datasetgen/mcts/schema.sql index ad9002813..b15956842 100644 --- a/src/datasetgen/mcts/schema.sql +++ b/src/datasetgen/mcts/schema.sql @@ -9,7 +9,9 @@ CREATE TABLE programs ( conversation_json JSONB, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, holdout_value: FLOAT DEFAULT 0.0, - raw_reward: FLOAT DEFAULT 0.0 + raw_reward: FLOAT DEFAULT 0.0, + version: INTEGER DEFAULT 0, -- version of the program code + version_description: TEXT DEFAULT NULL, -- description of the version change ); -- Task queue table diff --git a/src/factorio_rcon_utils.py b/src/factorio_rcon_utils.py index 347619ebb..c67b465e5 100644 --- a/src/factorio_rcon_utils.py +++ b/src/factorio_rcon_utils.py @@ -118,5 +118,9 @@ def _lua2python(command, response, *parameters, trace=False, start=0): if trace: print(f"failure: {command} \t") end = timer() - return lua.decode(response), (end - start) + + try: + return lua.decode(response), (end - start) + except Exception as e: + return None, (end - start) diff --git a/src/factorio_types.py b/src/factorio_types.py index ac3475731..563cdf0e7 100644 --- a/src/factorio_types.py +++ b/src/factorio_types.py @@ -54,7 +54,7 @@ class Prototype(enum.Enum): PumpJack = "pumpjack", PumpJack Boiler = "boiler", Boiler SteamEngine = "steam-engine", Generator - Pipe = "pipe", Entity + Pipe = "pipe", Pipe IronChest = "iron-chest", Chest WoodenChest = "wooden-chest", Chest IronGearWheel = "iron-gear-wheel", Entity diff --git a/src/llm_factory.py b/src/llm_factory.py index 1cff494df..ac32d2027 100644 --- a/src/llm_factory.py +++ b/src/llm_factory.py @@ -110,6 +110,7 @@ def call(self, *args, **kwargs): temperature=kwargs.get('temperature', 0.3), messages=kwargs.get('messages', None), logit_bias=kwargs.get('logit_bias', None), + n=kwargs.get('n', None), #stop=["\n\n"],#, "\n#"], #presence_penalty=1, #frequency_penalty=0.6, diff --git a/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md b/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md index 421686f22..e24799c52 100644 --- a/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md +++ b/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md @@ -18,4 +18,6 @@ Important things to remember - When you have something in your inventory, then do not make a task to create more of that something. Be diverse with tasks and inventories on the map - If you need to smelt anything, you need a stone furnace. If there are no stone furnaces on the map, you need to craft a stone furnace and you need 5 stone for that - When you put new items into a furnace to smelt them, you need to first extract the previous items in the furnace if there are any -- When you have placedan entity, you do not need to check if the placement was correct with assert is_close. The apiwill error out if the placement of an entity fails \ No newline at end of file +- When you have placed an entity, you do not need to check if the placement was correct with assert is_close. The api will error out if the placement of an entity fails. + +WHEN SETTING OBJECTIVES, DEVELOP AUTOMATED PRODUCTION CHAINS. \ No newline at end of file diff --git a/src/tests/actions/test_craft.py b/src/tests/actions/test_craft.py index 224e14527..3245459b1 100644 --- a/src/tests/actions/test_craft.py +++ b/src/tests/actions/test_craft.py @@ -8,7 +8,7 @@ def game(instance): instance.reset() instance.set_inventory(**{'iron-plate': 40, 'iron-gear-wheel': 1, - 'electronic-circuit': 2, + #'electronic-circuit': 2, 'pipe': 1, 'copper-plate': 10}) yield instance @@ -37,6 +37,10 @@ def test_craft_item(game): assert initial_iron_plate - final_iron_plate == iron_cost * quantity assert initial_iron_chest + quantity == final_iron_chest +def test_recursive_crafting(game): + crafted_circuits = game.craft_item(Prototype.ElectronicCircuit, quantity=4) + assert crafted_circuits + def test_craft_copper_coil(game): """ Craft 20 copper cable and verify that only 10 copper plates have been deducted. @@ -66,7 +70,7 @@ def test_craft_entity_with_missing_intermediate_resources(game): """ starting_stats = game.production_stats() # Craft 20 copper coil - crafted = game.craft_item(Prototype.OffshorePump, quantity=1) + crafted = game.craft_item(Prototype.ElectronicCircuit, quantity=1) # Check the production stats final_stats = game.production_stats() diff --git a/src/tests/complex/test_edge_cases.py b/src/tests/complex/test_edge_cases.py index e577494ab..4ce8d6d7a 100644 --- a/src/tests/complex/test_edge_cases.py +++ b/src/tests/complex/test_edge_cases.py @@ -163,8 +163,6 @@ def test_blueprint_functionality(game): assert any(e.prototype == Prototype.BurnerInserter for e in placed_entities.entities) assert any(e.prototype == Prototype.IronChest for e in placed_entities.entities) - - # Run the tests if __name__ == "__main__": pytest.main([__file__]) \ No newline at end of file From 87d375ac9be5d47d8901658b5f14eb642f1b7810 Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Thu, 21 Nov 2024 17:23:02 +0000 Subject: [PATCH 13/15] Add recursive crafting support. --- src/actions/craft_item.lua | 139 +++++++++++++++++++++++++++++++++++++ 1 file changed, 139 insertions(+) diff --git a/src/actions/craft_item.lua b/src/actions/craft_item.lua index 4cf5bc796..515e7ce8a 100644 --- a/src/actions/craft_item.lua +++ b/src/actions/craft_item.lua @@ -1,6 +1,145 @@ global.actions.craft_item = function(player_index, entity, count) local player = game.get_player(player_index) + -- Helper functions remain the same + local function get_missing_ingredients(player, recipe, count) + local missing_ingredients = {} + local crafts_needed = math.ceil(count / recipe.products[1].amount) + for _, ingredient in pairs(recipe.ingredients) do + local count_that_player_has = player.get_item_count(ingredient.name) + local needed = ingredient.amount * crafts_needed + if count_that_player_has < needed then + local difference = needed - count_that_player_has + missing_ingredients[ingredient.name] = difference + end + end + return missing_ingredients + end + + local function can_craft_recipe(player, recipe_name) + local recipe = player.force.recipes[recipe_name] + if not recipe then + return false, "recipe doesn't exist" + end + if not recipe.enabled then + return false, "recipe not unlocked" + end + if recipe.category ~= "crafting" then + return false, "recipe requires specific crafting machine" + end + return true, recipe + end + + local function update_production_stats(force, recipe, crafts_count) + local stats = force.item_production_statistics + for _, ingredient in pairs(recipe.ingredients) do + stats.on_flow(ingredient.name, -ingredient.amount * crafts_count) + end + for _, product in pairs(recipe.products) do + if product.type == "item" then + stats.on_flow(product.name, product.amount * crafts_count) + end + end + end + + -- Single recursive crafting function that handles both fast and slow modes + local function attempt_craft(player, entity_name, count, attempted_recipes) + attempted_recipes = attempted_recipes or {} + + -- Prevent infinite recursion + if attempted_recipes[entity_name] then + return 0, "recursive crafting loop detected" + end + attempted_recipes[entity_name] = true + + local can_craft, recipe_or_error = can_craft_recipe(player, entity_name) + if not can_craft then + return 0, recipe_or_error + end + + local recipe = recipe_or_error + local crafts_needed = math.ceil(count / recipe.products[1].amount) + local actual_craft_count = crafts_needed * recipe.products[1].amount + + -- Check for missing ingredients + local missing_ingredients = get_missing_ingredients(player, recipe, actual_craft_count) + if next(missing_ingredients) then + -- Try to craft each missing ingredient + for ingredient_name, needed_amount in pairs(missing_ingredients) do + local crafted_amount, error_msg = attempt_craft(player, ingredient_name, needed_amount, attempted_recipes) + if crafted_amount == 0 then + return 0, "couldn't craft intermediate " .. ingredient_name .. ": " .. error_msg + end + end + end + + -- After potentially crafting intermediates, check if we can now craft the original item + if global.fast then + -- Fast crafting implementation + local missing = get_missing_ingredients(player, recipe, actual_craft_count) + if next(missing) then + local missing_str = "" + for name, amount in pairs(missing) do + missing_str = missing_str .. name .. " x" .. amount .. ", " + end + return 0, "still missing ingredients: " .. missing_str:sub(1, -3) + end + + for _, ingredient in pairs(recipe.ingredients) do + player.remove_item({name = ingredient.name, count = ingredient.amount * crafts_needed}) + end + + local crafted = player.insert({name = entity_name, count = actual_craft_count}) + if crafted < actual_craft_count then + player.surface.spill_item_stack(player.position, {name = entity_name, count = actual_craft_count - crafted}) + end + + update_production_stats(player.force, recipe, crafted) + return crafted, nil + else + -- Slow crafting implementation + local crafted = player.begin_crafting{count=count, recipe=entity_name} + if crafted == 0 then + return 0, "unable to begin crafting - check prerequisites and inventory space" + end + update_production_stats(player.force, recipe, crafted) + return crafted, nil + end + end + + -- Main crafting logic + local total_crafted = 0 + local final_error = nil + + while total_crafted < count do + local remaining = count - total_crafted + local crafted_amount, error_msg = attempt_craft(player, entity, remaining, {}) + + if crafted_amount > 0 then + total_crafted = total_crafted + crafted_amount + if not global.fast then + break + end + else + final_error = error_msg + break + end + end + + if total_crafted >= count or (not global.fast and total_crafted > 0) then + return count + elseif total_crafted > 0 then + error(string.format("\"Successfully crafted %dx but failed to craft %dx %s because %s\"", + total_crafted, count - total_crafted, entity, final_error)) + else + error(string.format("\"Failed to craft %dx %s because %s\"", + count, entity, final_error)) + end +end + +global.actions.craft_item2 = function(player_index, entity, count) + local player = game.get_player(player_index) + -- Helper function to check missing ingredients local function get_missing_ingredients(player, recipe, count) local missing_ingredients = {} From f555448500f1eafcc87fec4df01ebe2155be7141 Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Fri, 22 Nov 2024 12:51:20 +0000 Subject: [PATCH 14/15] Add refined system policy. Fixed pipe connect when invalid connection points or jagged path --- cluster/local/create_docker_compose_config.py | 7 +- cluster/local/docker-compose-8.yml | 281 ++++++++++++++++++ src/actions/connect_entities.lua | 14 +- src/actions/get_resource_patch.lua | 2 +- src/actions/insert_item.lua | 2 +- src/actions/nearest.lua | 4 +- src/actions/place_entity_next_to.lua | 4 +- src/actions/place_object.lua | 4 +- .../auto_curriculum/plan_sampler.py | 12 +- src/datasetgen/mcts/__main__.py | 102 ++++--- src/datasetgen/mcts/mcts.py | 2 +- src/factorio_entities.py | 16 +- src/init/serialize.lua | 155 +++++++--- .../system_message_policy.md | 9 +- .../system_message_policy_refined.md | 79 +++++ .../system_message_policy_self_gen.md | 9 +- src/tests/actions/test_get_entities.py | 1 + src/tests/connect/test_connect_pipes.py | 44 ++- 18 files changed, 623 insertions(+), 124 deletions(-) create mode 100644 cluster/local/docker-compose-8.yml create mode 100644 src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_refined.md diff --git a/cluster/local/create_docker_compose_config.py b/cluster/local/create_docker_compose_config.py index b452b983f..70d534af1 100644 --- a/cluster/local/create_docker_compose_config.py +++ b/cluster/local/create_docker_compose_config.py @@ -29,6 +29,11 @@ def generate_compose_config(num_instances: int) -> Dict[str, Any]: "type": "bind", "source": "../scenarios/default_lab_scenario", "target": "/opt/factorio/scenarios/default_lab_scenario" + }, + { + "type": "bind", + "source": "~/Applications/Factorio.app/Contents/Resources/mods", + "target": "/opt/factorio/mods" } ], "ports": [ @@ -81,7 +86,7 @@ def setup_docker_compose(num_instances: int): time.sleep(10) if __name__ == "__main__": - num_instances = 4 + num_instances = 8 if len(sys.argv) != 2: print("Usage: python create_docker_compose_config.py ") else: diff --git a/cluster/local/docker-compose-8.yml b/cluster/local/docker-compose-8.yml new file mode 100644 index 000000000..110125e04 --- /dev/null +++ b/cluster/local/docker-compose-8.yml @@ -0,0 +1,281 @@ +services: + factorio_0: + command: /opt/factorio/bin/x64/factorio --start-server-load-scenario default_lab_scenario + --port 34197 --server-settings /opt/factorio/config/server-settings.json --map-gen-settings + /opt/factorio/config/map-gen-settings.json --map-settings /opt/factorio/config/map-settings.json + --server-banlist /opt/factorio/config/server-banlist.json --rcon-port 27015 + --rcon-password "factorio" --server-whitelist /opt/factorio/config/server-whitelist.json + --use-server-whitelist --server-adminlist /opt/factorio/config/server-adminlist.json + --mod-directory /opt/factorio/mods + deploy: + resources: + limits: + cpus: '1' + memory: 1024m + entrypoint: [] + environment: + - SAVES=/opt/factorio/saves + - CONFIG=/opt/factorio/config + - MODS=/opt/factorio/mods + - SCENARIOS=/opt/factorio/scenarios + - PORT=34197 + - RCON_PORT=27015 + image: factorio:latest + platform: linux/amd64 + ports: + - 34197:34197/udp + - 27015:27015/tcp + restart: unless-stopped + user: factorio + volumes: + - source: ../scenarios/default_lab_scenario + target: /opt/factorio/scenarios/default_lab_scenario + type: bind + - source: ~/Applications/Factorio.app/Contents/Resources/mods + target: /opt/factorio/mods + type: bind + factorio_1: + command: /opt/factorio/bin/x64/factorio --start-server-load-scenario default_lab_scenario + --port 34197 --server-settings /opt/factorio/config/server-settings.json --map-gen-settings + /opt/factorio/config/map-gen-settings.json --map-settings /opt/factorio/config/map-settings.json + --server-banlist /opt/factorio/config/server-banlist.json --rcon-port 27015 + --rcon-password "factorio" --server-whitelist /opt/factorio/config/server-whitelist.json + --use-server-whitelist --server-adminlist /opt/factorio/config/server-adminlist.json + --mod-directory /opt/factorio/mods + deploy: + resources: + limits: + cpus: '1' + memory: 1024m + entrypoint: [] + environment: + - SAVES=/opt/factorio/saves + - CONFIG=/opt/factorio/config + - MODS=/opt/factorio/mods + - SCENARIOS=/opt/factorio/scenarios + - PORT=34197 + - RCON_PORT=27015 + image: factorio:latest + platform: linux/amd64 + ports: + - 34198:34197/udp + - 27016:27015/tcp + restart: unless-stopped + user: factorio + volumes: + - source: ../scenarios/default_lab_scenario + target: /opt/factorio/scenarios/default_lab_scenario + type: bind + - source: ~/Applications/Factorio.app/Contents/Resources/mods + target: /opt/factorio/mods + type: bind + factorio_2: + command: /opt/factorio/bin/x64/factorio --start-server-load-scenario default_lab_scenario + --port 34197 --server-settings /opt/factorio/config/server-settings.json --map-gen-settings + /opt/factorio/config/map-gen-settings.json --map-settings /opt/factorio/config/map-settings.json + --server-banlist /opt/factorio/config/server-banlist.json --rcon-port 27015 + --rcon-password "factorio" --server-whitelist /opt/factorio/config/server-whitelist.json + --use-server-whitelist --server-adminlist /opt/factorio/config/server-adminlist.json + --mod-directory /opt/factorio/mods + deploy: + resources: + limits: + cpus: '1' + memory: 1024m + entrypoint: [] + environment: + - SAVES=/opt/factorio/saves + - CONFIG=/opt/factorio/config + - MODS=/opt/factorio/mods + - SCENARIOS=/opt/factorio/scenarios + - PORT=34197 + - RCON_PORT=27015 + image: factorio:latest + platform: linux/amd64 + ports: + - 34199:34197/udp + - 27017:27015/tcp + restart: unless-stopped + user: factorio + volumes: + - source: ../scenarios/default_lab_scenario + target: /opt/factorio/scenarios/default_lab_scenario + type: bind + - source: ~/Applications/Factorio.app/Contents/Resources/mods + target: /opt/factorio/mods + type: bind + factorio_3: + command: /opt/factorio/bin/x64/factorio --start-server-load-scenario default_lab_scenario + --port 34197 --server-settings /opt/factorio/config/server-settings.json --map-gen-settings + /opt/factorio/config/map-gen-settings.json --map-settings /opt/factorio/config/map-settings.json + --server-banlist /opt/factorio/config/server-banlist.json --rcon-port 27015 + --rcon-password "factorio" --server-whitelist /opt/factorio/config/server-whitelist.json + --use-server-whitelist --server-adminlist /opt/factorio/config/server-adminlist.json + --mod-directory /opt/factorio/mods + deploy: + resources: + limits: + cpus: '1' + memory: 1024m + entrypoint: [] + environment: + - SAVES=/opt/factorio/saves + - CONFIG=/opt/factorio/config + - MODS=/opt/factorio/mods + - SCENARIOS=/opt/factorio/scenarios + - PORT=34197 + - RCON_PORT=27015 + image: factorio:latest + platform: linux/amd64 + ports: + - 34200:34197/udp + - 27018:27015/tcp + restart: unless-stopped + user: factorio + volumes: + - source: ../scenarios/default_lab_scenario + target: /opt/factorio/scenarios/default_lab_scenario + type: bind + - source: ~/Applications/Factorio.app/Contents/Resources/mods + target: /opt/factorio/mods + type: bind + factorio_4: + command: /opt/factorio/bin/x64/factorio --start-server-load-scenario default_lab_scenario + --port 34197 --server-settings /opt/factorio/config/server-settings.json --map-gen-settings + /opt/factorio/config/map-gen-settings.json --map-settings /opt/factorio/config/map-settings.json + --server-banlist /opt/factorio/config/server-banlist.json --rcon-port 27015 + --rcon-password "factorio" --server-whitelist /opt/factorio/config/server-whitelist.json + --use-server-whitelist --server-adminlist /opt/factorio/config/server-adminlist.json + --mod-directory /opt/factorio/mods + deploy: + resources: + limits: + cpus: '1' + memory: 1024m + entrypoint: [] + environment: + - SAVES=/opt/factorio/saves + - CONFIG=/opt/factorio/config + - MODS=/opt/factorio/mods + - SCENARIOS=/opt/factorio/scenarios + - PORT=34197 + - RCON_PORT=27015 + image: factorio:latest + platform: linux/amd64 + ports: + - 34201:34197/udp + - 27019:27015/tcp + restart: unless-stopped + user: factorio + volumes: + - source: ../scenarios/default_lab_scenario + target: /opt/factorio/scenarios/default_lab_scenario + type: bind + - source: ~/Applications/Factorio.app/Contents/Resources/mods + target: /opt/factorio/mods + type: bind + factorio_5: + command: /opt/factorio/bin/x64/factorio --start-server-load-scenario default_lab_scenario + --port 34197 --server-settings /opt/factorio/config/server-settings.json --map-gen-settings + /opt/factorio/config/map-gen-settings.json --map-settings /opt/factorio/config/map-settings.json + --server-banlist /opt/factorio/config/server-banlist.json --rcon-port 27015 + --rcon-password "factorio" --server-whitelist /opt/factorio/config/server-whitelist.json + --use-server-whitelist --server-adminlist /opt/factorio/config/server-adminlist.json + --mod-directory /opt/factorio/mods + deploy: + resources: + limits: + cpus: '1' + memory: 1024m + entrypoint: [] + environment: + - SAVES=/opt/factorio/saves + - CONFIG=/opt/factorio/config + - MODS=/opt/factorio/mods + - SCENARIOS=/opt/factorio/scenarios + - PORT=34197 + - RCON_PORT=27015 + image: factorio:latest + platform: linux/amd64 + ports: + - 34202:34197/udp + - 27020:27015/tcp + restart: unless-stopped + user: factorio + volumes: + - source: ../scenarios/default_lab_scenario + target: /opt/factorio/scenarios/default_lab_scenario + type: bind + - source: ~/Applications/Factorio.app/Contents/Resources/mods + target: /opt/factorio/mods + type: bind + factorio_6: + command: /opt/factorio/bin/x64/factorio --start-server-load-scenario default_lab_scenario + --port 34197 --server-settings /opt/factorio/config/server-settings.json --map-gen-settings + /opt/factorio/config/map-gen-settings.json --map-settings /opt/factorio/config/map-settings.json + --server-banlist /opt/factorio/config/server-banlist.json --rcon-port 27015 + --rcon-password "factorio" --server-whitelist /opt/factorio/config/server-whitelist.json + --use-server-whitelist --server-adminlist /opt/factorio/config/server-adminlist.json + --mod-directory /opt/factorio/mods + deploy: + resources: + limits: + cpus: '1' + memory: 1024m + entrypoint: [] + environment: + - SAVES=/opt/factorio/saves + - CONFIG=/opt/factorio/config + - MODS=/opt/factorio/mods + - SCENARIOS=/opt/factorio/scenarios + - PORT=34197 + - RCON_PORT=27015 + image: factorio:latest + platform: linux/amd64 + ports: + - 34203:34197/udp + - 27021:27015/tcp + restart: unless-stopped + user: factorio + volumes: + - source: ../scenarios/default_lab_scenario + target: /opt/factorio/scenarios/default_lab_scenario + type: bind + - source: ~/Applications/Factorio.app/Contents/Resources/mods + target: /opt/factorio/mods + type: bind + factorio_7: + command: /opt/factorio/bin/x64/factorio --start-server-load-scenario default_lab_scenario + --port 34197 --server-settings /opt/factorio/config/server-settings.json --map-gen-settings + /opt/factorio/config/map-gen-settings.json --map-settings /opt/factorio/config/map-settings.json + --server-banlist /opt/factorio/config/server-banlist.json --rcon-port 27015 + --rcon-password "factorio" --server-whitelist /opt/factorio/config/server-whitelist.json + --use-server-whitelist --server-adminlist /opt/factorio/config/server-adminlist.json + --mod-directory /opt/factorio/mods + deploy: + resources: + limits: + cpus: '1' + memory: 1024m + entrypoint: [] + environment: + - SAVES=/opt/factorio/saves + - CONFIG=/opt/factorio/config + - MODS=/opt/factorio/mods + - SCENARIOS=/opt/factorio/scenarios + - PORT=34197 + - RCON_PORT=27015 + image: factorio:latest + platform: linux/amd64 + ports: + - 34204:34197/udp + - 27022:27015/tcp + restart: unless-stopped + user: factorio + volumes: + - source: ../scenarios/default_lab_scenario + target: /opt/factorio/scenarios/default_lab_scenario + type: bind + - source: ~/Applications/Factorio.app/Contents/Resources/mods + target: /opt/factorio/mods + type: bind diff --git a/src/actions/connect_entities.lua b/src/actions/connect_entities.lua index c2fd7e00a..d067faca8 100644 --- a/src/actions/connect_entities.lua +++ b/src/actions/connect_entities.lua @@ -29,13 +29,21 @@ local function are_fluidboxes_connected(entity1, entity2) end local function is_placeable(position) - --local tile = game.surfaces[1].get_tile(position.x, position.y) - --return tile.collides_with('player-layer') == false + -- Check if the tile is water or other impassable tiles + local invalid_tiles = { + ["water"] = true, + ["deepwater"] = true, + ["water-green"] = true, + ["deepwater-green"] = true, + ["water-shallow"] = true, + ["water-mud"] = true, + } + local entities = game.surfaces[1].find_entities_filtered{ position = position, collision_mask = "player-layer" } - return #entities == 0 + return #entities == 0 and not invalid_tiles[game.surfaces[1].get_tile(position.x, position.y).name] end local function find_placeable_neighbor(pos, previous_pos) diff --git a/src/actions/get_resource_patch.lua b/src/actions/get_resource_patch.lua index 34c6fdc6f..53702ee6c 100644 --- a/src/actions/get_resource_patch.lua +++ b/src/actions/get_resource_patch.lua @@ -52,7 +52,7 @@ global.actions.get_resource_patch = function(player_index, resource, x, y, radiu else local resource_entities = surface.find_entities_filtered{position = position, name = resource, radius = radius} if #resource_entities == 0 then - error("No resource of type " .. resource .. " at the specified location.") + error("\"No resource of type " .. resource .. " at the specified location.\"") end -- Recursive function to explore all connected resource entities diff --git a/src/actions/insert_item.lua b/src/actions/insert_item.lua index d61f128a6..c9d64e1fa 100644 --- a/src/actions/insert_item.lua +++ b/src/actions/insert_item.lua @@ -6,7 +6,7 @@ global.actions.insert_item = function(player_index, insert_item, count, x, y) -- Check if player has enough items local item_count = player.get_item_count(insert_item) if item_count == 0 then - error('No '..insert_item..' to place') + error('\"No '..insert_item..' to place\"') end local closest_distance = math.huge diff --git a/src/actions/nearest.lua b/src/actions/nearest.lua index b9775abd3..fb5ea7f5a 100644 --- a/src/actions/nearest.lua +++ b/src/actions/nearest.lua @@ -40,7 +40,7 @@ global.actions.nearest = function(player_index, resource) end end if closest == nil then - error("Could not find an entity called "..resource) + error("\"Could not find an entity called "..resource.."\"") end --return { x = position.x - closest.x, y = position.y - closest.y } return {x = closest.x, y = closest.y} @@ -60,7 +60,7 @@ global.actions.nearest = function(player_index, resource) end if closest == nil then - error("Could not find an entity called "..resource) + error("\"Could not find an entity called "..resource.."\"") end game.print("Distance to "..resource.." is "..closest_distance) return {x= closest.x, y = closest.y} diff --git a/src/actions/place_entity_next_to.lua b/src/actions/place_entity_next_to.lua index 9d5cc7d3e..8b9224bfc 100644 --- a/src/actions/place_entity_next_to.lua +++ b/src/actions/place_entity_next_to.lua @@ -247,8 +247,8 @@ global.actions.place_entity_next_to = function(player_index, entity, ref_x, ref_ game.print(e.type) table.insert(entity_names, e.name) end - error("Cannot place entity at the position " .. serpent.line(new_position) .. " with direction " .. - serpent.line(orientation) .. ". Attempting to place next to: "..ref_entity.name..". Nearby entities: " .. serpent.line(entity_names)) + error("\'Cannot place entity at the position " .. serpent.line(new_position) .. " with direction " .. + serpent.line(orientation) .. ". Attempting to place next to: "..ref_entity.name..". Nearby entities: " .. serpent.line(entity_names).."\'") end local new_entity = player.surface.create_entity({ diff --git a/src/actions/place_object.lua b/src/actions/place_object.lua index 7cb9f92c1..ba5540751 100644 --- a/src/actions/place_object.lua +++ b/src/actions/place_object.lua @@ -141,7 +141,7 @@ global.actions.place_entity = function(player_index, entity, direction, x, y, ex return global.actions.get_entity(player_index, entity, new_position.x, new_position.y) end else - error("Could not find a suitable position to place " .. entity .. " near the target location.") + error("\"Could not find a suitable position to place " .. entity .. " near the target location.\"") end else -- Clear existing entities if exact placement is required @@ -318,7 +318,7 @@ global.actions.place_entity2 = function(player_index, entity, direction, x, y, e return global.actions.get_entity(player_index, entity, new_position.x, new_position.y) end else - error("Could not find a suitable position to place " .. entity .. " near the target location.") + error("\"Could not find a suitable position to place " .. entity .. " near the target location.\"") end else --local existing_entity = global.actions.get_entity(player_index, entity, position.x, position.y) diff --git a/src/datasetgen/auto_curriculum/plan_sampler.py b/src/datasetgen/auto_curriculum/plan_sampler.py index c271c5212..2be96ad51 100644 --- a/src/datasetgen/auto_curriculum/plan_sampler.py +++ b/src/datasetgen/auto_curriculum/plan_sampler.py @@ -60,21 +60,23 @@ def get_unsupervised_objective(self, instance): self.init_system_prompt(instance) mining_setup = self.get_mining_setup(instance) messages = [{"role": "system", "content": self.system_prompt}] - user_message = f"Your starting inventory is {starting_inventory}. Your initial mining setup is: {mining_setup}. Create a useful task that you can carry out in the current game and the python script to achieve the task" + user_message = f"Your starting inventory is {starting_inventory}. Your initial mining setup is: {mining_setup}. In light of previous instructions, create a useful task that you can carry out in the current game and the python script to achieve the task" user_message += f"\n{self.planning_addition_for_prompt}" messages.append({"role": "user", "content": user_message}) response = self.llm_factory.call(messages=messages, model = self.model, temperature=0.7, - max_tokens=256, + max_tokens=512, stop_sequences = ["```"]) - - full_output = response.choices[0].message.content + try: + full_output = response.choices[0].message.content + except: + full_output = response.content[0].text.strip() full_output = full_output.strip() new_line_idx = full_output.rfind("\n") full_output = full_output[:new_line_idx].replace("Sure!", "").strip() - full_output = f'"""\n{full_output}\n"""' + #full_output = f'"""\n{full_output}\n"""' return full_output, response def get_mining_setup(self, instance): diff --git a/src/datasetgen/mcts/__main__.py b/src/datasetgen/mcts/__main__.py index 313585b69..d90ad0934 100644 --- a/src/datasetgen/mcts/__main__.py +++ b/src/datasetgen/mcts/__main__.py @@ -1,5 +1,6 @@ import json import os +import random import statistics from datasetgen.auto_curriculum.plan_sampler import PlanSampler @@ -68,56 +69,59 @@ async def get_seed_programs( scenarios = [f for f in os.listdir(plan_sampler.starting_scenarios_folder) if os.path.isdir(os.path.join(plan_sampler.starting_scenarios_folder, f))] - seeds_per_scenario = n_seeds // len(scenarios) - remaining_seeds = n_seeds % len(scenarios) - - for scenario in scenarios: - num_samples = seeds_per_scenario + (1 if remaining_seeds > 0 else 0) - remaining_seeds -= 1 - - for _ in range(num_samples): - game_state = plan_sampler.get_game_state(instance, scenario) - if not game_state: - continue - - objective, response = plan_sampler(instance, game_state) - - if len(objective) < 100: - continue - - - - conversation = Conversation(messages=[ - Message(role="system", content=mcts.system_prompt), - Message(role="user", - content=f"Inventory: {json.dumps(game_state.inventory.__dict__)}\n\n{PLANNING_ADDITION_PROMPT}"), - Message(role="assistant", content=objective) - ]) - messages = conversation.model_dump()['messages'] - if not objective.strip().startswith('"""'): - objective = '"""\n'+objective - - program = Program( - id=hash((objective, json.dumps(messages))), - code=objective, - conversation=conversation, - value=default_reward, - state=game_state, - version=mcts.version, - version_description=mcts.version_description, - token_usage=response.usage.total_tokens if hasattr(response, 'usage') else None, - completion_token_usage=response.usage.completion_tokens if hasattr(response, 'usage') else None, - prompt_token_usage=response.usage.prompt_tokens if hasattr(response, 'usage') else None, - ) - seeded_programs.append(program) + # Select a random N number of scenarios + selected_scenarios = random.sample(scenarios, n_seeds) + + for scenario in selected_scenarios: + game_state = plan_sampler.get_game_state(instance, scenario) + if not game_state: + continue + + objective, response = plan_sampler(instance, game_state) + + if len(objective) < 100: + continue + + conversation = Conversation(messages=[ + Message(role="system", content=mcts.system_prompt), + Message(role="user", + content=f"Starting Inventory: {json.dumps(game_state.inventory.__dict__)}"), + Message(role="assistant", content=objective) + ]) + messages = conversation.model_dump()['messages'] + if not objective.strip().startswith('"""'): + objective = '"""\n'+objective + + try: + token_usage = response.usage.total_tokens if hasattr(response, 'usage') else None + completion_token_usage = response.usage.completion_tokens if hasattr(response, 'usage') else None + prompt_token_usage = response.usage.prompt_tokens if hasattr(response, 'usage') else None + except: + completion_token_usage = response.usage.output_tokens + prompt_token_usage = response.usage.input_tokens + token_usage = prompt_token_usage + completion_token_usage + + program = Program( + id=hash((objective, json.dumps(messages))), + code=objective, + conversation=conversation, + value=default_reward, + state=game_state, + version=mcts.version, + version_description=mcts.version_description, + token_usage=token_usage, + completion_token_usage=completion_token_usage, + prompt_token_usage=prompt_token_usage, + ) + seeded_programs.append(program) return seeded_programs async def main(): - model = "ft:gpt-4o-2024-08-06:paperplane-ai:fact-self-gen-planning:AQzcPI91" - prompt_path = "../../prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md" - version = 8 - version_description = "Seeded / No planning prompt in user messages / Step-wise evaluation / Errors not saved" + model = "ft:gpt-4o-2024-08-06:paperplane-ai:fact-self-gen-planning:AQzcPI91"#"o1-mini" #"gpt-4o" #"ft:gpt-4o-2024-08-06:paperplane-ai:fact-self-gen-planning:AQzcPI91" + prompt_path = "../../prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_refined.md" + version = 12 + version_description = "Seeded / Base model / No planning prompt in user messages / Step-wise evaluation / Refined system prompt" # Initialize components llm = LLMFactory(model) @@ -139,7 +143,7 @@ async def main(): # Get execution directory from __file__ or other source execution_dir = os.path.dirname(os.path.realpath(__file__)) # load from prompts/bottoms_up_prompts/system_message_policy_self_gen into string - with open("../../prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md", "r") as f: + with open("../../prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_refined.md", "r") as f: system_prompt = f.read().format(schema=instances[0].get_system_prompt()) print("Initializing MCTS...") @@ -168,7 +172,7 @@ async def main(): sampler = PlanSampler(model, prompt_path, starting_scenario_folder) print("Sampling seed scenarios...") - seeded_programs = await get_seed_programs(mcts, sampler, n_seeds=0) + seeded_programs = await get_seed_programs(mcts, sampler, n_seeds=3) for program in seeded_programs: await db_client.create_program(program) @@ -176,7 +180,7 @@ async def main(): best_programs = await mcts.search( n_iterations=500, samples_per_iteration=len(instances)-1, # One for each instance, minus a holdout. - skip_failures=False, + skip_failures=True, ) print("\nBest programs found:") diff --git a/src/datasetgen/mcts/mcts.py b/src/datasetgen/mcts/mcts.py index 869928c69..13267cace 100644 --- a/src/datasetgen/mcts/mcts.py +++ b/src/datasetgen/mcts/mcts.py @@ -90,7 +90,7 @@ async def _generate_programs_batch(self, conversation: Conversation, n_samples: messages=formatted_messages, max_tokens=2048, n=n_samples, - temperature=0.7, # Adjust as needed + temperature=1, logit_bias=logit_bias ) diff --git a/src/factorio_entities.py b/src/factorio_entities.py index 6d81a8f78..365792203 100644 --- a/src/factorio_entities.py +++ b/src/factorio_entities.py @@ -327,6 +327,9 @@ class Lab(Entity): lab_input: Inventory = Inventory() lab_modules: Inventory = Inventory() +class Pipe(Entity): + pass + class EntityGroup(BaseModel): input_positions: List[Position] position: Position @@ -336,8 +339,13 @@ class BeltGroup(EntityGroup): belts: List[TransportBelt] output_positions: List[Position] -class Pipe(Entity): - #connections: List[Position] = [] - pass + def __repr__(self) -> str: + belt_summary = f"[{len(self.belts)} belts]" + return f"BeltGroup(position={self.position}, input_positions={self.input_positions}, output_positions={self.output_positions}, status={self.status}, belts={belt_summary})" + class PipeGroup(EntityGroup): - pipes: List[Pipe] \ No newline at end of file + pipes: List[Pipe] + + def __repr__(self) -> str: + pipe_summary = f"[{len(self.pipes)} pipes]" + return f"PipeGroup(position={self.position}, input_positions={self.input_positions}, status={self.status}, pipes={pipe_summary})" \ No newline at end of file diff --git a/src/init/serialize.lua b/src/init/serialize.lua index 2a7d22ed3..5d2381d35 100644 --- a/src/init/serialize.lua +++ b/src/init/serialize.lua @@ -390,6 +390,7 @@ global.utils.serialize_fluidbox = function(fluidbox) return serialized end + local function get_offshore_pump_pipe_position(entity) local x, y = entity.position.x, entity.position.y local orientation = entity.orientation * 8 @@ -506,25 +507,6 @@ function get_boiler_pipe_positions(entity) return pipe_positions end -function add_burner_inventory2(burner) - local fuel_inventory = burner.inventory - if fuel_inventory and #fuel_inventory > 0 then - local serialized = {} - for i = 1, #fuel_inventory do - local item = fuel_inventory[i] - if item and item.valid_for_read then - local item_name = "\"" .. item.name .. "\"" - if serialized[item_name] then - serialized[item_name] = serialized[item_name] + item.count - else - serialized[item_name] = item.count - end - end - end - return serialized - end - return {} -end function add_burner_inventory(serialized, burner) local fuel_inventory = burner.inventory @@ -653,24 +635,87 @@ function get_inverse_entity_direction(entity, factorio_direction) else return -1 end - --elseif prototype and prototype.type == "mining-drill" then - -- if factorio_direction == defines.direction.east then - -- return defines.direction.east - -- elseif factorio_direction == defines.direction.south then - -- return defines.direction.south - -- elseif factorio_direction == defines.direction.west then - -- return defines.direction.west - -- else -- north - -- return defines.direction.north - -- end else game.print("Returning direction: " .. math.floor(factorio_direction / 2) .. ', '.. factorio_direction) -- For other entity types, convert Factorio's direction to 0-3 range return factorio_direction - --return factorio_direction end end +-- Helper function to check if a position is valid (not colliding with water or other impassable tiles) +local function is_valid_connection_point(surface, position) + -- Get the tile at the position + local tile = surface.get_tile(position.x, position.y) + + -- Check if the tile is water or other impassable tiles + local invalid_tiles = { + ["water"] = true, + ["deepwater"] = true, + ["water-green"] = true, + ["deepwater-green"] = true, + ["water-shallow"] = true, + ["water-mud"] = true, + } + + -- Return false if the tile is invalid, true otherwise + return not invalid_tiles[tile.name] +end + +-- Helper function to filter connection points +local function filter_connection_points(entity, points) + if not points then return nil end + + local filtered_points = {} + for _, point in ipairs(points) do + if is_valid_connection_point(entity.surface, point) then + table.insert(filtered_points, point) + end + end + + -- If all points were filtered out, return nil + if #filtered_points == 0 then + return nil + end + + return filtered_points +end + +-- Modified pipe position functions to include filtering +local function get_pipe_positions_filtered(entity) + local positions = get_pipe_positions(entity) + return filter_connection_points(entity, positions) +end + +local function get_pumpjack_pipe_position_filtered(entity) + local positions = get_pumpjack_pipe_position(entity) + return filter_connection_points(entity, positions) +end + +local function get_boiler_pipe_positions_filtered(entity) + local positions = get_boiler_pipe_positions(entity) + + -- Special handling for boiler since it has a different structure + if not positions then return nil end + + local filtered = { + water_inputs = filter_connection_points(entity, positions.water_inputs), + steam_output = positions.steam_output -- Usually steam output doesn't need filtering as it connects to pipes above ground + } + + -- If all water inputs were filtered out, return nil + if not filtered.water_inputs or #filtered.water_inputs == 0 then + return nil + end + + return filtered +end + +local function get_offshore_pump_pipe_position_filtered(entity) + local positions = get_offshore_pump_pipe_position(entity) + return filter_connection_points(entity, positions) +end + + global.entity_status_names = { [defines.entity_status.working] = "working", [defines.entity_status.normal] = "normal", @@ -728,12 +773,6 @@ global.utils.serialize_entity = function(entity) --game.print("Serializing entity: " .. entity.name .. " with direction: " .. entity.direction) local direction = entity.direction - -- This is needed because the entity.direction on the map is not always the actual direction - -- (e.g inserters and offshore pumps have opposite directions on the map to the actual cardinal) - --if entity.direction ~= nil then - -- direction = get_inverse_entity_direction(entity.name, entity.direction/2)*2--get_inverse_entity_direction(entity.name, entity.direction) - --end - if direction ~= nil then direction = get_entity_direction(entity.name, entity.direction) else @@ -1046,19 +1085,14 @@ global.utils.serialize_entity = function(entity) serialized.connection_points = {{x = x + 0.5, y = y - 2}, {x = x + 0.5, y = y + 2}} serialized.steam_output_point = {x = x - 2, y = y} end - - --serialized.fluid_input_point = entity.fluidbox.get_connections(1)[1].position end if entity.type == "generator" then - serialized.connection_points = get_pipe_positions(entity) - - --create_beam_point(game.players[1], serialized.connection_points[1]) - --create_beam_point(game.players[1], serialized.connection_points[2]) + serialized.connection_points = get_pipe_positions_filtered(entity) end if entity.name == "pumpjack" then - serialized.connection_points = get_pumpjack_pipe_position(entity) + serialized.connection_points = get_pumpjack_pipe_position_filtered(entity) end -- Add fuel and input ingredients if the entity is a furnace or burner @@ -1100,6 +1134,41 @@ global.utils.serialize_entity = function(entity) serialized.direction = get_inverse_entity_direction(entity.name, entity.direction) --api_direction_map[entity.direction] + -- Post-process connection points if they exist + if serialized.connection_points then + local filtered_points = {} + for _, point in ipairs(serialized.connection_points) do + if is_valid_connection_point(game.surfaces[1], point) then + table.insert(filtered_points, point) + end + end + + -- Update connection points or remove if all were filtered + if #filtered_points > 0 then + serialized.connection_points = filtered_points + else + serialized.connection_points = nil + end + + -- Add warning if points were filtered + if not serialized.warnings then + serialized.warnings = {} + end + if #filtered_points < #serialized.connection_points then + table.insert(serialized.warnings, "Some connection points were filtered due to distance from entity center") + end + end + + -- Handle special case for boilers which have separate steam output points + if serialized.steam_output_point then + if not is_valid_connection_point(game.surfaces[1], serialized.steam_output_point) then + serialized.steam_output_point = nil + if not serialized.warnings then + serialized.warnings = {} + end + table.insert(serialized.warnings, "Steam output point was filtered due to distance from entity center") + end + end return serialized end diff --git a/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy.md b/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy.md index 034ad4315..2829dfdac 100644 --- a/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy.md +++ b/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy.md @@ -1,7 +1,6 @@ -You are an AI agent creating Python policy scripts that carry out actions in the Factorio game to achieve a objective. YOu are given a objective and must create the python policy using the factorio api to carry out that objective. The policy must be one script that defines all the steps required to achieve the objective. The script you write will define steps that are required to be carried out to successfully achieve the objective. Before each step you first think what is the next step you need to make in python comments. You then write the code interacting with the game API that carries out this step. You also write assert statements after your steps to ensure the steps were carried out correctly. You must test thoroughly to ensure steps were carried out correctly. Make sure to test the final outcome of the policy with asserts to ensure the objective that was given has been achieved. -The API for factorio you need to use is the following: +You are an AI agent exploring the world of Factotio. -You have access to the following Game API for use in your Python code: -{schema} +You create Python scripts to interact with the game and to achieve a objective. Invent appropriate objectives, and write Python to achieve them. The script you write will define steps that are required to be carried out to successfully achieve the objective. Before each step, you first think what is the next step using comments. You then write the code to carry out this step. Ensure the step succeeds by writing assert statements. Observe the game entities and reason about how they can fit together to generate flows of resources. -You are supplied with the game state and inventory and the objective by the user. Output any comments you have and a full python script that carries out the objective between ```python and ``` tags \ No newline at end of file +You have access to the following Game API for use in your Python code: +{schema} \ No newline at end of file diff --git a/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_refined.md b/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_refined.md new file mode 100644 index 000000000..d47b7b0d7 --- /dev/null +++ b/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_refined.md @@ -0,0 +1,79 @@ +# Factorio Factory Building Agent + +You are an AI agent tasked with developing increasingly sophisticated automated production systems in Factorio. Your primary goal is to design and implement efficient resource chains and automated factories using Python scripts. + +## Core Objectives +- Design and implement automated production chains +- Optimize resource gathering and processing +- Create scalable factory layouts +- Develop interconnected production systems +- Progress through technology tiers systematically + +## Planning Approach +For each objective you tackle: +1. Analyze available resources and current factory state +2. Plan the production chain components needed +3. Consider resource dependencies and bottlenecks +4. Design the physical layout and connections +5. Implement error checking and validation + +## Implementation Guidelines +Your Python scripts should: +- Use detailed comments to explain the reasoning behind each step +- Include assert statements to validate operations +- Consider resource efficiency and throughput +- Build upon existing infrastructure when possible +- Implement proper error handling + +## Technical Requirements +When writing code: +1. Entity Placement Rules: + - Move to position before placing entities (`move_to` required) + - Connect entities using transport belts via `connect_entities(source.drop_position, target.pickup_position)` + - Burner inserters required for item transfer between entities + +2. Resource Management: + - Check furnace fuel levels using `furnace.fuel.get('Prototype.Coal')` + - Verify entity contents using `inspect_inventory()` + - Extract items from furnaces before new smelting operations + +3. Crafting Requirements: + - Print the recipes for things you want to craft first + +4. Optimization Rules: + - Avoid duplicating existing resources + - Prioritize using available resources before crafting new ones + - Consider resource chain dependencies + +5. Observation + - Use `print` statements or logging for debugging + - Observe entities with `get_entities()` + - Observe inventories with `inspect_inventory()` + +6. Score + - Your score is based on resource generation throughput from your automated system + - Maximise your score + +## Available Game API +{schema} + +## Output Format +Provide your solution as a complete Python script between triple backticks: +```python +# Your implementation here +``` + +## Factory Development Guidelines +Your factory designs should demonstrate: +1. Scalability - Allow for future expansion +2. Efficiency - Minimize resource waste +3. Automation - Reduce manual intervention +4. Integration - Connect different production systems +5. Progressive Complexity - Build towards more advanced products +6. Fix warnings - Warnings in observed entities indicate that it requires fixing (e.g missing ingredients) + +Focus on creating sustainable, automated production chains that can be expanded and integrated into larger systems. Each new objective should build upon previous infrastructure with the aim of maximising throughput. + +Remember to document your thinking process as Python comments, and explain how each new addition contributes to the overall factory. + +If an error is returned from your actions, you should reflect on why your actions caused this error, and attempt to fix it by observing the game state and submitting a fix. \ No newline at end of file diff --git a/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md b/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md index e24799c52..a1038bc18 100644 --- a/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md +++ b/src/prompts/bottoms_up_prompts/finetuning_prompts/system_message_policy_self_gen.md @@ -1,12 +1,13 @@ -You are an AI agent creating Python policy scripts that carry out actions in the Factorio game to achieve a objective. You must first create a task that is achievable and useful in the current factorio game state and then must create the python policy using the factorio api to carry out that objective. The policy must be one script that defines all the steps required to achieve the objective. The script you write will define steps that are required to be carried out to successfully achieve the objective. Before each step you first think what is the next step you need to make in python comments. You then write the code interacting with the game API that carries out this step. You also write assert statements after your steps to ensure the steps were carried out correctly. You must test thoroughly to ensure steps were carried out correctly. Make sure to test the final outcome of the policy with asserts to ensure the objective that was given has been achieved. -The API for factorio you need to use is the following: +You are an AI agent exploring the world of Factorio. + +You create Python scripts to interact with the game and to achieve your objectives. Invent appropriate objectives, and write Python to achieve them. The script you write will define steps that are required to be carried out to successfully achieve the objective. Before each step, you first think what is the next step using comments. You then write the code to carry out this step. Ensure the step succeeds by writing assert statements. Observe the game entities and reason about how they can fit together to generate flows of resources. You have access to the following Game API for use in your Python code: {schema} -You are supplied with the game state and inventory by the user. Output the useful objective that can be carried out, any planning you need to do and then the python script that carries out the objective. The full python script that carries out the objective must be between ```python and ``` tags +The full python script that carries out the objective must be between ```python and ``` tags. -Important things to remember +Important things to remember: - To put items into a entity you need a burner inserter, for instance to put items into a chest, you need to put a burner inserter next to the chest and rotate it - To connect different items and entities with each other, you need to connect them with transport belts with connect_entities command, for instance connect_entities(drill.drop_position, inserter.pickup_position) - Always take into account the starting inventory and any resources you can use from there or from chests or furnaces on the map diff --git a/src/tests/actions/test_get_entities.py b/src/tests/actions/test_get_entities.py index b1cc0370b..7636315c5 100644 --- a/src/tests/actions/test_get_entities.py +++ b/src/tests/actions/test_get_entities.py @@ -99,6 +99,7 @@ def test_get_contiguous_transport_belts(game): transport_belts = game.get_entities({Prototype.TransportBelt}, start_position) + transport_belt_string = str(transport_belts) assert len(transport_belts) == 1, "Failed to retrieve transport belts" def test_get_filtered_entities(game): # put down a chest at origin diff --git a/src/tests/connect/test_connect_pipes.py b/src/tests/connect/test_connect_pipes.py index e0dd22055..f1d07a433 100644 --- a/src/tests/connect/test_connect_pipes.py +++ b/src/tests/connect/test_connect_pipes.py @@ -3,7 +3,7 @@ import pytest -from factorio_entities import Entity, Position, PipeGroup, EntityStatus +from factorio_entities import Entity, Position, PipeGroup, EntityStatus, ResourcePatch from factorio_instance import Direction from factorio_types import Prototype, Resource, PrototypeName @@ -250,3 +250,45 @@ def test_avoid_self_collision(game): assert len(pipes) == 1, "Failed to construct a single contiguous pipe group" assert pipes, "Failed to connect offshore pump to boiler with pipes" print("Successfully connected offshore pump to boiler with pipes") + +def test_connect_where_connection_points_are_blocked(game): + # Move to the water source + water_source = game.nearest(Resource.Water) + game.move_to(water_source) + print(f"Moved to water source at {water_source}") + # Place the offshore pump + pump = game.place_entity(Prototype.OffshorePump, Direction.UP, water_source) + print(f"Placed offshore pump at {pump.position}") + """ + Step 2: Place the boiler and connect it to the pump + """ + # Calculate position for the boiler (4 tiles away from the pump) + boiler_position = Position(x=pump.position.x + 4, y=pump.position.y) + + # Move to the calculated position + game.move_to(boiler_position) + print(f"Moved to boiler position at {boiler_position}") + + # Place the boiler + boiler = game.place_entity(Prototype.Boiler, Direction.UP, boiler_position) + print(f"Placed boiler at {boiler.position}") + + # Connect the pump to the boiler with pipes + pump_to_boiler_pipes = game.connect_entities(pump, boiler, Prototype.Pipe) + assert pump_to_boiler_pipes, "Failed to connect pump to boiler with pipes" + print("Connected pump to boiler with pipes") + + assert boiler.connection_points[0] in pump_to_boiler_pipes[0].input_positions + +def test_connect_ragged_edges(game): + water: ResourcePatch = game.get_resource_patch(Resource.Water, game.nearest(Resource.Water)) + + start_pos = water.bounding_box.left_top + end_pos = water.bounding_box.right_bottom + + # Move to the start position and place an offshore pump + game.move_to(start_pos) + + pipes = game.connect_entities(start_pos, end_pos, Prototype.Pipe) + + assert len(pipes) == 1 \ No newline at end of file From 02cdb07067829b584bccf9f681ad132ed49c73ae Mon Sep 17 00:00:00 2001 From: Jack Hopkins Date: Fri, 22 Nov 2024 12:52:51 +0000 Subject: [PATCH 15/15] Removing connection points after setting them in serialized --- src/init/serialize.lua | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/init/serialize.lua b/src/init/serialize.lua index 5d2381d35..0e8cfc2bc 100644 --- a/src/init/serialize.lua +++ b/src/init/serialize.lua @@ -1088,11 +1088,11 @@ global.utils.serialize_entity = function(entity) end if entity.type == "generator" then - serialized.connection_points = get_pipe_positions_filtered(entity) + serialized.connection_points = get_pipe_positions(entity) end if entity.name == "pumpjack" then - serialized.connection_points = get_pumpjack_pipe_position_filtered(entity) + serialized.connection_points = get_pumpjack_pipe_position(entity) end -- Add fuel and input ingredients if the entity is a furnace or burner