diff --git a/src_code_for_reproducibility/__pycache__/__init__.cpython-311.pyc b/src_code_for_reproducibility/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b817647d28cfd5c3fd8ea2da23692ec7b1aaacc Binary files /dev/null and b/src_code_for_reproducibility/__pycache__/__init__.cpython-311.pyc differ diff --git a/src_code_for_reproducibility/chat_utils/apply_template.py b/src_code_for_reproducibility/chat_utils/apply_template.py new file mode 100644 index 0000000000000000000000000000000000000000..f143ab15e6b30637700df3381b7de686153017cd --- /dev/null +++ b/src_code_for_reproducibility/chat_utils/apply_template.py @@ -0,0 +1,78 @@ +import torch + +from mllm.chat_utils.chat_turn import ChatTurn +from mllm.chat_utils.template_specific import ( + custom_llama3_template, + custom_qwen2_template, + custom_qwen3_template, + qwen2_assistant_postfix, + qwen3_assistant_postfix, +) + + +def get_custom_chat_template(tokenizer) -> str: + """ + Get the chat template for the tokenizer. + """ + if "qwen2" in tokenizer.name_or_path.lower(): + return custom_qwen2_template + elif "llama" in tokenizer.name_or_path.lower(): + return custom_llama3_template + elif "qwen3" in tokenizer.name_or_path.lower(): + return custom_qwen3_template + else: + raise ValueError(f"Tokenizer {tokenizer.name_or_path} not supported") + + +def get_custom_assistant_postfix(tokenizer) -> torch.Tensor: + """ + Get the custom assistant postfix for the tokenizer. + """ + if "qwen2" in tokenizer.name_or_path.lower(): + return qwen2_assistant_postfix + elif "qwen3" in tokenizer.name_or_path.lower(): + return qwen3_assistant_postfix + return torch.tensor([], dtype=torch.long) + + +def tokenize_chats(chats: list[ChatTurn], tokenizer, enable_thinking) -> None: + """ + Set the chat_template_token_ids for each chat turn. + # TODO: use engine tokens if available + """ + custom_template = get_custom_chat_template(tokenizer) + custom_assistant_postfix: torch.Tensor = get_custom_assistant_postfix(tokenizer) + for i, chat in enumerate(chats): + if chat.chat_template_token_ids is None: + if chat.role == "user": + next_chat = chats[i + 1] if i + 1 < len(chats) else None + add_generation_prompt = True + if next_chat and next_chat.role == "user": + add_generation_prompt = False + encoded_chat = tokenizer.apply_chat_template( + [chat], + return_tensors="pt", + chat_template=custom_template, + add_generation_prompt=add_generation_prompt, + add_system_prompt=True if i == 0 else False, + enable_thinking=enable_thinking, + ).flatten() + previous_chat = chats[i - 1] if i > 0 else None + if previous_chat and previous_chat.role == "assistant": + encoded_chat = torch.cat([custom_assistant_postfix, encoded_chat]) + elif chat.role == "assistant": + encoded_chat = chat.out_token_ids + chat.chat_template_token_ids = encoded_chat + + +def chat_turns_to_token_ids( + chats: list[ChatTurn], tokenizer, enable_thinking +) -> list[int]: + """ + Tokenize the chat turns and set the chat_template_token_ids for each chat turn. + """ + tokenize_chats(chats=chats, tokenizer=tokenizer, enable_thinking=enable_thinking) + token_ids = [] + for chat in chats: + token_ids.append(chat.chat_template_token_ids) + return torch.cat(token_ids) diff --git a/src_code_for_reproducibility/chat_utils/chat_turn.py b/src_code_for_reproducibility/chat_utils/chat_turn.py new file mode 100644 index 0000000000000000000000000000000000000000..d22c1b853def374fa0ee21eb523f2b9d104ed35b --- /dev/null +++ b/src_code_for_reproducibility/chat_utils/chat_turn.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, List, Literal, Optional, Tuple + +import jsonschema +import torch +from pydantic import BaseModel, ConfigDict, Field, model_validator + +AgentId = str + + +class ChatTurn(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) # needed for torch tensors + + role: str = Field(pattern="^(user|assistant)$") + agent_id: AgentId # ID of the agent with which the chat occured + content: str + reasoning_content: str | None = None + chat_template_token_ids: torch.LongTensor | None = None # Token ids of chat template format. For example, token ids of "{content}"" + out_token_ids: torch.LongTensor | None = ( + None # tokens generated from inference engine + ) + log_probs: torch.FloatTensor | None = None + is_state_end: bool = False # indicates whether this chat turn marks the end of a state in the trajectory diff --git a/src_code_for_reproducibility/chat_utils/template_specific.py b/src_code_for_reproducibility/chat_utils/template_specific.py new file mode 100644 index 0000000000000000000000000000000000000000..5d8f8e5f1be6bcafd1b9e1a465c446253b985e3e --- /dev/null +++ b/src_code_for_reproducibility/chat_utils/template_specific.py @@ -0,0 +1,87 @@ +import huggingface_hub +import torch +from transformers import AutoTokenizer + +custom_llama3_template = """ +{%- if add_system_prompt %} + {{- '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nCutting Knowledge Date: December 2023\nToday Date: 26 Jul 2024\n\n<|eot_id|>' }} +{%- endif %} +{%- for message in messages %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }} +{%- endfor %} + +{%- if add_generation_prompt %} + {{- '<|start_header_id|>' + 'assistant' + '<|end_header_id|>\n\n' }} +{%- endif %} +""" + +qwen2_assistant_postfix = ( + AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct") + .encode("\n", return_tensors="pt") + .flatten() +) +qwen3_assistant_postfix = ( + AutoTokenizer.from_pretrained("Qwen/Qwen3-8B") + .encode("\n", return_tensors="pt") + .flatten() +) +custom_qwen2_template = """ +{%- if add_system_prompt %} + {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if loop.index0 > ns.last_query_index %} + {%- if reasoning_content %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} +{%- endif %} +""" + +custom_qwen3_template = """ +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- endif %} +{%- endif %} +""" diff --git a/src_code_for_reproducibility/docs/Makefile b/src_code_for_reproducibility/docs/Makefile new file mode 100644 index 0000000000000000000000000000000000000000..9375655022388557d65910465b09d2a5fa9d2e4a --- /dev/null +++ b/src_code_for_reproducibility/docs/Makefile @@ -0,0 +1,19 @@ +# Minimal makefile for Sphinx documentation + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS) diff --git a/src_code_for_reproducibility/docs/generate_docs.py b/src_code_for_reproducibility/docs/generate_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..e644cbbf091a97420500fb47346c07be5ed141ac --- /dev/null +++ b/src_code_for_reproducibility/docs/generate_docs.py @@ -0,0 +1,249 @@ +#!/usr/bin/env python3 +""" +Script to automatically generate Sphinx documentation for all modules and build the HTML website. +""" +import importlib.util +import os +import subprocess +import sys + + +def check_and_install_dependencies(): + """Check for required dependencies and install them if missing.""" + required_packages = [ + "sphinx", + "sphinx-rtd-theme", + "sphinxcontrib-napoleon", + "sphinxcontrib-mermaid", + "sphinx-autodoc-typehints", + ] + + missing_packages = [] + + for package in required_packages: + # Convert package name to module name (replace - with _) + module_name = package.replace("-", "_") + + # Check if the package is installed + if importlib.util.find_spec(module_name) is None: + missing_packages.append(package) + + # Install missing packages + if missing_packages: + print(f"Installing missing dependencies: {', '.join(missing_packages)}") + subprocess.check_call( + [sys.executable, "-m", "pip", "install"] + missing_packages + ) + print("Dependencies installed successfully") + else: + print("All required dependencies are already installed") + + +def create_makefile(docs_dir): + """Create a Makefile for Sphinx documentation if it doesn't exist.""" + makefile_path = os.path.join(docs_dir, "Makefile") + + if os.path.exists(makefile_path): + print(f"Makefile already exists at {makefile_path}") + return + + print(f"Creating Makefile at {makefile_path}") + + makefile_content = """# Minimal makefile for Sphinx documentation + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS) +""" + + with open(makefile_path, "w") as f: + f.write(makefile_content) + + print("Makefile created successfully") + + +def create_make_bat(docs_dir): + """Create a make.bat file for Windows if it doesn't exist.""" + make_bat_path = os.path.join(docs_dir, "make.bat") + + if os.path.exists(make_bat_path): + print(f"make.bat already exists at {make_bat_path}") + return + + print(f"Creating make.bat at {make_bat_path}") + + make_bat_content = """@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd +""" + + with open(make_bat_path, "w") as f: + f.write(make_bat_content) + + print("make.bat created successfully") + + +def main(): + # Check and install required dependencies + print("=== Checking dependencies ===") + check_and_install_dependencies() + + # Get the directory of this script + script_dir = os.path.dirname(os.path.abspath(__file__)) + + # Path to the project root + project_root = os.path.dirname(script_dir) + + # Path to the source directory + source_dir = os.path.join(project_root, "src") + + # Path to the docs source directory + docs_source_dir = os.path.join(script_dir, "source") + + # Print paths for debugging + print(f"Script directory: {script_dir}") + print(f"Project root: {project_root}") + print(f"Source directory: {source_dir}") + print(f"Docs source directory: {docs_source_dir}") + + # Make sure the source directory exists + if not os.path.exists(source_dir): + print(f"Error: Source directory {source_dir} does not exist!") + sys.exit(1) + + # Make sure the docs source directory exists + if not os.path.exists(docs_source_dir): + print(f"Creating docs source directory: {docs_source_dir}") + os.makedirs(docs_source_dir) + + # Step 1: Run sphinx-apidoc to generate .rst files for all modules + print("\n=== Generating API documentation ===") + cmd = [ + "sphinx-apidoc", + "-f", # Force overwriting of existing files + "-e", # Put module documentation before submodule documentation + "-M", # Put module documentation before subpackage documentation + "-o", + docs_source_dir, # Output directory + source_dir, # Source code directory + ] + + print(f"Running command: {' '.join(cmd)}") + result = subprocess.run(cmd, capture_output=True, text=True) + + # Print the output of the command + print("STDOUT:") + print(result.stdout) + + print("STDERR:") + print(result.stderr) + + if result.returncode != 0: + print(f"Error: sphinx-apidoc failed with return code {result.returncode}") + sys.exit(1) + + # List the files in the docs source directory + print("\nFiles in docs/source directory:") + for file in sorted(os.listdir(docs_source_dir)): + print(f" {file}") + + print("\nDocumentation source files generated successfully!") + + # Step 2: Create Makefile and make.bat if they don't exist + create_makefile(script_dir) + create_make_bat(script_dir) + + # Step 3: Build the HTML documentation + print("\n=== Building HTML documentation ===") + + # Determine the build command based on the platform + if os.name == "nt": # Windows + build_cmd = ["make.bat", "html"] + else: # Unix/Linux/Mac + build_cmd = ["make", "html"] + + # Change to the docs directory to run the build command + os.chdir(script_dir) + + print(f"Running command: {' '.join(build_cmd)}") + build_result = subprocess.run(build_cmd, capture_output=True, text=True) + + # Print the output of the build command + print("STDOUT:") + print(build_result.stdout) + + print("STDERR:") + print(build_result.stderr) + + if build_result.returncode != 0: + print(f"Error: HTML build failed with return code {build_result.returncode}") + sys.exit(1) + + # Get the path to the built HTML documentation + html_dir = os.path.join(script_dir, "build", "html") + index_path = os.path.join(html_dir, "index.html") + + if os.path.exists(index_path): + print(f"\nHTML documentation built successfully!") + print(f"You can view it by opening: {index_path}") + + # Try to open the documentation in a browser + try: + import webbrowser + + print("\nAttempting to open documentation in your default browser...") + webbrowser.open(f"file://{index_path}") + except Exception as e: + print(f"Could not open browser automatically: {e}") + else: + print(f"\nWarning: HTML index file not found at {index_path}") + + +if __name__ == "__main__": + main() diff --git a/src_code_for_reproducibility/docs/make.bat b/src_code_for_reproducibility/docs/make.bat new file mode 100644 index 0000000000000000000000000000000000000000..dc1312ab09ca6fb0267dee6b28a38e69c253631a --- /dev/null +++ b/src_code_for_reproducibility/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/src_code_for_reproducibility/markov_games/__init__.py b/src_code_for_reproducibility/markov_games/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src_code_for_reproducibility/markov_games/agent.py b/src_code_for_reproducibility/markov_games/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..9961d860fecca7eaa4d13cb195b50dc658be5bec --- /dev/null +++ b/src_code_for_reproducibility/markov_games/agent.py @@ -0,0 +1,76 @@ +""" +In simple RL paradise, where the action dimensions are constant and well defined, +Agent classes are not necessary. But in MARL, with LLM's, there isn't always +a direct path from policy to action. For instance, from the observation of the environment, +a prompt must be created. Then, the outputs of the policy might be incorrect, so a second +request to the LLM must be sent before the action is well defined. This is why this Agent class exists. +It acts as a mini environment, bridging the gap between the core simulation and +the LLM policies. +""" + +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any, Tuple + +from numpy.random import default_rng + +from mllm.markov_games.rollout_tree import AgentActLog + + +class Agent(ABC): + @abstractmethod + def __init__( + self, + seed: int, + agent_id: str, + agent_name: str, + agent_policy: Callable[[list[dict]], str], + *args, + **kwargs, + ): + """ + Initialize the agent state. + """ + self.seed = seed + self.agent_id = agent_id + self.agent_name = agent_name + self.policy = policy + self.rng = default_rng(self.seed) + raise NotImplementedError + + async def act(self, observation) -> Tuple[Any, AgentActLog]: + """ + Query (possibly multiple times) a policy (or possibly a pool of policies) to + obtain the action of the agent. + + Example: + action = None + prompt = self.observation_to_prompt(observation) + while not self.valid(action): + output = await self.policy.generate(prompt) + action = self.policy_output_to_action(output) + return action + + Returns: + action + step_info + """ + raise NotImplementedError + + def get_safe_copy(self): + """ + Return copy of the agent object that is decorrelated from the original object. + """ + raise NotImplementedError + + def reset(self): + raise NotImplementedError + + def render(self): + raise NotImplementedError + + def close(self): + raise NotImplementedError + + def get_agent_info(self): + raise NotImplementedError diff --git a/src_code_for_reproducibility/markov_games/alternative_actions_runner.py b/src_code_for_reproducibility/markov_games/alternative_actions_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..c64db2deda539a1a71e045309cfdf257d2cbc614 --- /dev/null +++ b/src_code_for_reproducibility/markov_games/alternative_actions_runner.py @@ -0,0 +1,138 @@ +import asyncio +import copy +import json +import os.path +from typing import Any, Tuple + +from mllm.markov_games.markov_game import AgentAndActionSafeCopy, MarkovGame +from mllm.markov_games.rollout_tree import ( + AgentActLog, + RolloutTreeBranchNode, + RolloutTreeNode, + RolloutTreeRootNode, + StepLog, +) + +AgentId = str + + + +async def run_with_unilateral_alt_action( + markov_game: MarkovGame, + agent_id: AgentId, + time_step: int, + branch_node: RolloutTreeBranchNode, + max_depth: int, +): + """ + This function is used to generate a new branch for a given agent. + """ + + # Generate alternative action and take a step + await markov_game.set_action_of_agent(agent_id) + terminated: bool = markov_game.take_simulation_step() + step_log = markov_game.get_step_log() + first_alternative_node = RolloutTreeNode( + step_log=step_log, + time_step=time_step, + ) + + # Generate rest of trajectory up to max depth + time_step += 1 + counter = 1 + previous_node = first_alternative_node + while not terminated and counter <= max_depth: + terminated, step_log = await markov_game.step() + current_node = RolloutTreeNode(step_log=step_log, time_step=time_step) + previous_node.child = current_node + previous_node = current_node + counter += 1 + time_step += 1 + + if branch_node.branches == None: + branch_node.branches = {agent_id: [first_alternative_node]} + else: + agent_branches = branch_node.branches.get(agent_id, []) + agent_branches.append(first_alternative_node) + branch_node.branches[agent_id] = agent_branches + + +async def AlternativeActionsRunner( + markov_game: MarkovGame, + output_folder: str, + nb_alternative_actions: int, + max_depth: int, + branch_only_on_new_round: bool = False, +): + """ + This method generates a trajectory with partially completed branches, + where the branching comes from taking unilateraly different actions. + The resulting data is used to estimate the updated advantage alignment policy gradient terms. + Let k := nb_sub_steps. Then the number of steps generated is O(Tk), where T is + the maximum trajectory length. + """ + + tasks = [] + time_step = 0 + terminated = False + root = RolloutTreeRootNode( + id=markov_game.get_id(), + crn_id=markov_game.get_crn_id() + ) + previous_node = root + + while not terminated: + mg_before_action = markov_game.get_safe_copy() + + # Get safe copies for main branch + agent_action_safe_copies: dict[ + AgentId, AgentAndActionSafeCopy + ] = await markov_game.get_actions_of_agents_without_side_effects() + + markov_game.set_actions_of_agents_manually(agent_action_safe_copies) + terminated = markov_game.take_simulation_step() + main_node = RolloutTreeNode( + step_log=markov_game.get_step_log(), time_step=time_step + ) + branch_node = RolloutTreeBranchNode(main_child=main_node) + previous_node.child = branch_node + previous_node = main_node + + # Get alternative branches by generating new unilateral actions + for agent_id in markov_game.agent_ids: + for _ in range(nb_alternative_actions): + # Get safe copies for branches + branch_agent_action_safe_copies: dict[ + AgentId, AgentAndActionSafeCopy + ] = { + agent_id: AgentAndActionSafeCopy( + action=copy.deepcopy(agent_action_safe_copy.action), + action_info=copy.deepcopy(agent_action_safe_copy.action_info), + agent_after_action=agent_action_safe_copy.agent_after_action.get_safe_copy(), + ) + for agent_id, agent_action_safe_copy in agent_action_safe_copies.items() + } + mg_branch: MarkovGame = mg_before_action.get_safe_copy() + other_agent_id = [id for id in mg_branch.agent_ids if id != agent_id][0] + mg_branch.set_action_and_agent_after_action_manually( + agent_id=other_agent_id, + agent_action_safe_copy=branch_agent_action_safe_copies[ + other_agent_id + ], + ) + task = asyncio.create_task( + run_with_unilateral_alt_action( + markov_game=mg_branch, + time_step=time_step, + agent_id=agent_id, + branch_node=branch_node, + max_depth=max_depth, + ) + ) + tasks.append(task) + time_step += 1 + + # wait for all branches to complete + await asyncio.gather(*tasks) + + return root diff --git a/src_code_for_reproducibility/markov_games/group_timesteps.py b/src_code_for_reproducibility/markov_games/group_timesteps.py new file mode 100644 index 0000000000000000000000000000000000000000..dad5271c500f539f2719110bd676f183746e51e4 --- /dev/null +++ b/src_code_for_reproducibility/markov_games/group_timesteps.py @@ -0,0 +1,150 @@ +""" +This module contains the logic for grouping time steps. +""" +import copy +from typing import Callable + +from mllm.markov_games.markov_game import MarkovGame +from mllm.markov_games.rollout_tree import ( + AgentActLog, + RolloutTreeBranchNode, + RolloutTreeNode, + RolloutTreeRootNode, + StepLog, +) +from mllm.markov_games.simulation import SimulationStepLog + +AgentId = str + + +def group_time_steps( + rollout_tree: RolloutTreeRootNode, + accumulation_stop_condition: Callable[[StepLog], bool], +) -> RolloutTreeRootNode: + """ + During generation, we create rollout trees according to the real time steps. + However, during training, we might want to treat groups of time steps as a single time step. + As a concrete example, take Trust-and-Split. At each round, say we have X time steps of communication and then one time step for the split. + Then the communication actions will not get any reward, and the split action will get the reward. During REINFORCE training, with discounting, this + can cause training instability. We could instead treat every action in the round as being part of a single action, and give it the reward of the split action. + This method helps to do this sort of grouping. + It accumulates actions until the accumulation_stop_condition is met, and then creates a new node with the accumulated actions. + It then recursively calls itself on the child node. + Details: + - The reward for the group is the reward of the last time step in the group. + - The simulation log for the group is the simulation log of the last time step in the group. + - The state end for the group becomes the first state end in the group. + - The agent info for the group is the agent info of the last time step in the group. + """ + + def group_step_logs(step_logs: list[StepLog]) -> StepLog: + """ + Concatenate per-agent chat turns across steps; keep only the first is_state_end. + """ + last_sim_log = step_logs[-1].simulation_step_log + agent_ids = {aid for s in step_logs for aid in s.action_logs.keys()} + grouped_logs: dict[AgentId, AgentActLog] = {} + for aid in agent_ids: + turns = [] + for s in step_logs: + act = s.action_logs.get(aid) + if act and act.chat_turns: + turns.extend(copy.deepcopy(act.chat_turns)) + disable_is_state_end = False + # Only the first state_end should be True, the rest should be False + for t in turns: + if t.is_state_end: + if disable_is_state_end: + t.is_state_end = False + else: + disable_is_state_end = True + continue + grouped_logs[aid] = AgentActLog( + chat_turns=turns, info=step_logs[-1].action_logs[aid].info + ) + return StepLog(action_logs=grouped_logs, simulation_step_log=last_sim_log) + + def group_time_steps_rec( + current_node: RolloutTreeNode | RolloutTreeBranchNode, + group_time_step: int, + accumulation_step_logs: list[StepLog], + ) -> RolloutTreeNode | RolloutTreeBranchNode: + """ + Groups time steps. Recursion is used to handle branches. + """ + assert isinstance(current_node, RolloutTreeNode) or isinstance( + current_node, RolloutTreeBranchNode + ), "Current node must be a tree node or a branch node. Is of type: " + str( + type(current_node) + ) + first_group_node = None + current_group_node = None + while current_node is not None: + if isinstance(current_node, RolloutTreeBranchNode): + raise Exception( + "Grouping timesteps by round is not supported for branching trajectories yet." + ) + # Special recursive case for branches + # if isinstance(current_node, RolloutTreeBranchNode): + # branches = {} + # for agent_id, branch_nodes in current_node.branches.items(): + # branch_group_nodes = [] + # for branch_node in branch_nodes: + # branch_group_node = group_time_steps_rec( + # current_node=branch_node, + # group_time_step=group_time_step, + # accumulation_step_logs=copy.deepcopy(accumulation_step_logs)) + # branch_group_nodes.append(branch_group_node) + # branches[agent_id] = branch_group_nodes + + # main_child_group_node = group_time_steps_rec( + # current_node=current_node.main_child, + # group_time_step=group_time_step, + # accumulation_step_logs=copy.deepcopy(accumulation_step_logs)) + + # return RolloutTreeBranchNode(main_child=main_child_group_node, branches=branches) + + # Accumulate + accumulation_step_logs.append(current_node.step_log) + if accumulation_stop_condition(current_node.step_log): + grouped_step_logs = group_step_logs(accumulation_step_logs) + accumulation_step_logs = [] + new_group_node = RolloutTreeNode( + step_log=grouped_step_logs, time_step=group_time_step, child=None + ) + if first_group_node == None: + first_group_node = new_group_node + group_time_step += 1 + if current_group_node is not None: + current_group_node.child = new_group_node + current_group_node = new_group_node + current_node = current_node.child + return first_group_node + + node = group_time_steps_rec( + current_node=rollout_tree.child, group_time_step=0, accumulation_step_logs=[] + ) + return RolloutTreeRootNode( + id=rollout_tree.id, + crn_id=rollout_tree.crn_id, + child=node, + agent_ids=rollout_tree.agent_ids, + ) + + +def stop_when_round_ends(step_log: StepLog) -> bool: + """ + Simplest stop condition. Will return True if step log is the last time step of a round. + This will throw an error if this information is not available in the simulation info. + """ + assert ( + "is_last_timestep_in_round" in step_log.simulation_step_log.info.keys() + ), "To group by round, is_last_timestep_in_round must be set in the info of your simulation step log at each time step." + return step_log.simulation_step_log.info["is_last_timestep_in_round"] + + +def group_by_round(rollout_tree: RolloutTreeRootNode) -> RolloutTreeRootNode: + """ + Groups time steps by round. + """ + return group_time_steps(rollout_tree, stop_when_round_ends) diff --git a/src_code_for_reproducibility/markov_games/linear_runner.py b/src_code_for_reproducibility/markov_games/linear_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..f81ab6d37ae9c680b2f2a53388988117d37f8a47 --- /dev/null +++ b/src_code_for_reproducibility/markov_games/linear_runner.py @@ -0,0 +1,30 @@ +import asyncio +import json +import os.path + +from mllm.markov_games.markov_game import MarkovGame +from mllm.markov_games.rollout_tree import RolloutTreeNode, RolloutTreeRootNode + + +async def LinearRunner( + markov_game: MarkovGame, output_folder: str +) -> RolloutTreeRootNode: + """ + This method generates a trajectory without branching. + """ + time_step = 0 + terminated = False + root = RolloutTreeRootNode( + id=markov_game.get_id(), + crn_id=markov_game.get_crn_id(), + agent_ids=markov_game.get_agent_ids(), + ) + previous_node = root + while not terminated: + terminated, step_log = await markov_game.step() + current_node = RolloutTreeNode(step_log=step_log, time_step=time_step) + previous_node.child = current_node + previous_node = current_node + time_step += 1 + + return root diff --git a/src_code_for_reproducibility/markov_games/markov_game.py b/src_code_for_reproducibility/markov_games/markov_game.py new file mode 100644 index 0000000000000000000000000000000000000000..73a48213bddcf0a59976fa0870eec19f59ae47d9 --- /dev/null +++ b/src_code_for_reproducibility/markov_games/markov_game.py @@ -0,0 +1,208 @@ +""" +This class unifies a simulation, and the agents acting in it (see `simulation.py` & `agent.py`). +In a MarkovGame step, + 1) each agent takes an action, + 2) the state transitions with respect to these actions, + 3) all relevant data of the step is appended to the historical data list + +In order to perform 3), the agents and the simulation are expected, at each time step, +to return a log of the state transition (from their perspective). +For instance, the Simulation might send rewards and the agents might send prompting contexts to be used later to generate the training data. +A different approach would be to simply have the agents keep their data private and log it upon completion of a trajectory. +The approach we use here centralizes the data gathering aspect, +making it easy to create sub-trajectories (in the `runners` defined in `runners.py`) descriptions that +only log information for step transitions occuring after the branching out. +""" +import asyncio +import copy +import json +import os +from dataclasses import dataclass +from typing import Any, List, Literal, Optional, Tuple + +from transformers.models.idefics2 import Idefics2Config + +from mllm.markov_games.agent import Agent +from mllm.markov_games.rollout_tree import AgentActLog, StepLog +from mllm.markov_games.simulation import Simulation + +AgentId = str + + +@dataclass +class AgentAndActionSafeCopy: + action: Any + action_info: AgentActLog + agent_after_action: type[Agent] + + +class MarkovGame(object): + def __init__( + self, + id: int, + agents: dict[AgentId, type[Agent]], + simulation: type[Simulation], + crn_id: int, + ): + """ + Args: + agents: + output_path: + Path where the step infos are saved. + simulation: + Simulation object. Example: IPDSimulation + """ + self.agents = agents + self.agent_ids = self.agents.keys() + self.simulation = simulation + self.simulation_step_log = None + self.agent_step_logs = {agent_id: None for agent_id in self.agent_ids} + self.actions = {} + self.id = id + self.crn_id = crn_id + + def get_id(self) -> str: + return self.id + + def get_crn_id(self) -> int: + return self.crn_id + + def get_agent_ids(self) -> List[AgentId]: + return list(self.agent_ids) + + async def get_action_of_agent_without_side_effects( + self, agent_id: AgentId + ) -> Tuple[Any, AgentActLog]: + """ + Safe function to get an action of an agent without modifying the agent or the simulation. + """ + agent = self.agents[agent_id] + agent_before_action = agent.get_safe_copy() + obs = self.simulation.get_obs_agent(agent_id) + action, action_info = await agent.act(observation=obs) + self.agents[agent_id] = agent_before_action + agent_after_action = agent.get_safe_copy() + return AgentAndActionSafeCopy(action, action_info, agent_after_action) + + async def get_actions_of_agents_without_side_effects( + self, + ) -> dict[AgentId, AgentAndActionSafeCopy]: + """ + Safe function to get an action of an agent without modifying the agent or the simulation. + """ + tasks = [] + for agent_id in self.agent_ids: + task = asyncio.create_task( + self.get_action_of_agent_without_side_effects(agent_id) + ) + tasks.append(task) + agent_and_action_safe_copies: list[ + AgentAndActionSafeCopy + ] = await asyncio.gather(*tasks) + return { + agent_id: agent_and_action_safe_copy + for agent_id, agent_and_action_safe_copy in zip( + self.agent_ids, agent_and_action_safe_copies + ) + } + + def set_action_and_agent_after_action_manually( + self, + agent_id: AgentId, + agent_action_safe_copy: AgentAndActionSafeCopy, + ): + """ + Set the action and the agent after action manually. + """ + self.actions[agent_id] = agent_action_safe_copy.action + self.agent_step_logs[agent_id] = agent_action_safe_copy.action_info + self.agents[agent_id] = agent_action_safe_copy.agent_after_action + + def set_actions_of_agents_manually( + self, actions: dict[AgentId, AgentAndActionSafeCopy] + ): + """ + Set the actions of agents manually. + """ + for agent_id, agent_action_safe_copy in actions.items(): + self.set_action_and_agent_after_action_manually( + agent_id, agent_action_safe_copy + ) + + async def set_action_of_agent(self, agent_id: AgentId): + """ + TOWRITE + """ + agent = self.agents[agent_id] + obs = self.simulation.get_obs_agent(agent_id) + action, action_info = await agent.act(observation=obs) + self.actions[agent_id] = action + self.agent_step_logs[agent_id] = action_info + + async def set_actions(self): + """ + TOWRITE + """ + # background_tasks = set() + tasks = [] + for agent_id in self.agent_ids: + task = asyncio.create_task(self.set_action_of_agent(agent_id)) + tasks.append(task) + await asyncio.gather(*tasks) + + def take_simulation_step(self): + """ + TOWRITE + """ + terminated, self.simulation_step_log = self.simulation.step(self.actions) + return terminated + + def get_step_log(self) -> StepLog: + """ + TOWRITE + TODO: assert actions and simulation have taken step + """ + step_log = StepLog( + simulation_step_log=self.simulation_step_log, + action_logs=self.agent_step_logs, + ) + return step_log + + async def step(self) -> Tuple[bool, StepLog]: + """ + TOWRITE + """ + await self.set_actions() + terminated = self.take_simulation_step() + step_log = self.get_step_log() + return terminated, step_log + + def get_safe_copy(self): + """ + TOWRITE + """ + + new_markov_game = copy.copy(self) + new_simulation = self.simulation.get_safe_copy() + new_agents = { + agent_id: agent.get_safe_copy() for agent_id, agent in self.agents.items() + } + + # Reassign copied components + new_markov_game.simulation = new_simulation + new_markov_game.agents = new_agents + + # IMPORTANT: ensure agent_ids references the new agents dict, not the original + new_markov_game.agent_ids = new_markov_game.agents.keys() + + # Deep-copy step data to avoid correlation + new_markov_game.simulation_step_log = copy.deepcopy(self.simulation_step_log) + new_markov_game.actions = copy.deepcopy(self.actions) + # Rebuild logs to align exactly with new agent ids + old_agent_step_logs = copy.deepcopy(self.agent_step_logs) + new_markov_game.agent_step_logs = { + agent_id: old_agent_step_logs.get(agent_id) + for agent_id in new_markov_game.agent_ids + } + + return new_markov_game diff --git a/src_code_for_reproducibility/markov_games/mg_utils.py b/src_code_for_reproducibility/markov_games/mg_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..27e0eb606ea6cfd5682c3b21a53e783c23e9b43f --- /dev/null +++ b/src_code_for_reproducibility/markov_games/mg_utils.py @@ -0,0 +1,89 @@ +import asyncio +import copy +from collections.abc import Callable +from dataclasses import dataclass + +from mllm.markov_games.ipd.ipd_agent import IPDAgent +from mllm.markov_games.ipd.ipd_simulation import IPD +from mllm.markov_games.markov_game import MarkovGame +from mllm.markov_games.negotiation.dond_agent import DealNoDealAgent +from mllm.markov_games.negotiation.dond_simulation import DealNoDealSimulation +from mllm.markov_games.negotiation.nego_hard_coded_policies import ( + HardCodedNegoGreedyPolicy, + HardCodedNegoWelfareMaximizingPolicy, +) +from mllm.markov_games.ipd.Ipd_hard_coded_agents import AlwaysCooperateIPDAgent, AlwaysDefectIPDAgent +from mllm.markov_games.negotiation.no_press_nego_agent import NoPressAgent +from mllm.markov_games.negotiation.no_press_nego_simulation import NoPressSimulation +from mllm.markov_games.negotiation.tas_agent import TrustAndSplitAgent +from mllm.markov_games.negotiation.tas_rps_agent import TrustAndSplitRPSAgent +from mllm.markov_games.negotiation.tas_rps_simulation import TrustAndSplitRPSSimulation +from mllm.markov_games.negotiation.tas_simple_agent import TrustAndSplitSimpleAgent +from mllm.markov_games.negotiation.tas_simple_simulation import ( + TrustAndSplitSimpleSimulation, +) +from mllm.markov_games.negotiation.tas_simulation import TrustAndSplitSimulation +from mllm.markov_games.rollout_tree import ( + AgentActLog, + RolloutTreeBranchNode, + RolloutTreeNode, + RolloutTreeRootNode, + StepLog, +) +from mllm.markov_games.simulation import SimulationStepLog + +AgentId = str + + +@dataclass +class AgentConfig: + agent_id: str + agent_name: str + agent_class_name: str + policy_id: str + init_kwargs: dict + + +@dataclass +class MarkovGameConfig: + id: int + seed: int + simulation_class_name: str + simulation_init_args: dict + agent_configs: list[AgentConfig] + + +def init_markov_game_components( + config: MarkovGameConfig, policies: dict[str, Callable[[list[dict]], str]] +): + """ + TOWRITE + """ + agents = {} + agent_names = [] + for agent_config in config.agent_configs: + agent_id = agent_config.agent_id + agent_name = agent_config.agent_name + agent_class = eval(agent_config.agent_class_name) + agent = agent_class( + seed=config.seed, + agent_id=agent_id, + agent_name=agent_name, + policy=policies[agent_config.policy_id], + **agent_config.init_kwargs, + ) + agents[agent_id] = agent + agent_names.append(agent_name) + simulation = eval(config.simulation_class_name)( + seed=config.seed, + agent_ids=list(agents.keys()), + agent_names=agent_names, + **config.simulation_init_args, + ) + markov_game = MarkovGame( + id=config.id, + crn_id=config.seed, + agents=agents, + simulation=simulation, + ) + return markov_game diff --git a/src_code_for_reproducibility/markov_games/rollout_tree.py b/src_code_for_reproducibility/markov_games/rollout_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..5e90fa6d7a0e0b644ccaf3533a1513568d21868f --- /dev/null +++ b/src_code_for_reproducibility/markov_games/rollout_tree.py @@ -0,0 +1,86 @@ +""" +TODO: add parent to nodes so that some verification can be done. For instance, to ensure that node reward keys match the parent node. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, List, Literal, Optional, Tuple + +import jsonschema +from pydantic import BaseModel, Field, model_validator + +from mllm.chat_utils.chat_turn import ChatTurn + +AgentId = str + + +class SimulationStepLog(BaseModel): + rewards: dict[AgentId, float] + info: Any = None + + +class AgentActLog(BaseModel): + chat_turns: list[ChatTurn] | None + info: Any = None + + @model_validator(mode="after") + def _exactly_one_state_end(self): + """ + This method is used to enforce that for each AgentActLog, there is exactly one ChatTurn which is a state end. + """ + if self.chat_turns != []: + n = sum(1 for t in self.chat_turns if t.is_state_end) + if n != 1: + raise ValueError( + f"AgentActLog must have exactly one ChatTurn with is_state_end=True; got {self.chat_turns}." + ) + return self + else: + return self + + +class StepLog(BaseModel): + action_logs: dict[AgentId, AgentActLog] + simulation_step_log: SimulationStepLog + + +# BranchType = Literal["unilateral_deviation", "common_deviation"] # might not be necessary +# class BranchNodeInfo(BaseModel): +# branch_id: str +# branch_for: AgentId +# branch_type: BranchType + + +class RolloutTreeNode(BaseModel): + step_log: StepLog + time_step: int + child: RolloutTreeNode | RolloutTreeBranchNode | None = None + + +class RolloutTreeBranchNode(BaseModel): + """ + First item of the tuple indicates which agent "called" for an alternative branch. + """ + + main_child: RolloutTreeNode + branches: dict[AgentId, list[RolloutTreeNode]] | None = None + + +class RolloutTreeRootNode(BaseModel): + id: int + crn_id: int # ID of the rng used to generate this rollout tree + child: RolloutTreeNode | RolloutTreeBranchNode | None = None + agent_ids: List[AgentId] = Field(min_length=1) + + +# class RolloutTreeLeafNode(BaseModel): +# step_log: StepLog +# time_step: int + + +# Necessary for self-referential stuff in pydantic +RolloutTreeBranchNode.model_rebuild() +RolloutTreeNode.model_rebuild() diff --git a/src_code_for_reproducibility/markov_games/run_markov_games.py b/src_code_for_reproducibility/markov_games/run_markov_games.py new file mode 100644 index 0000000000000000000000000000000000000000..08b84024668ed4375453d1fd515a78eb9bb23414 --- /dev/null +++ b/src_code_for_reproducibility/markov_games/run_markov_games.py @@ -0,0 +1,24 @@ +import asyncio +from collections.abc import Callable +from dataclasses import dataclass + +from torch._C import ClassType + +from mllm.markov_games.markov_game import MarkovGame +from mllm.markov_games.rollout_tree import RolloutTreeRootNode + + +async def run_markov_games( + runner: Callable[[MarkovGame], RolloutTreeRootNode], + runner_kwargs: dict, + output_folder: str, + markov_games: list[MarkovGame], +) -> list[RolloutTreeRootNode]: + tasks = [] + for mg in markov_games: + tasks.append( + asyncio.create_task( + runner(markov_game=mg, output_folder=output_folder, **runner_kwargs) + ) + ) + return await asyncio.gather(*tasks) diff --git a/src_code_for_reproducibility/markov_games/simulation.py b/src_code_for_reproducibility/markov_games/simulation.py new file mode 100644 index 0000000000000000000000000000000000000000..b0b804e2aa4c288b3d98cc8106cfd727f1cc1e1a --- /dev/null +++ b/src_code_for_reproducibility/markov_games/simulation.py @@ -0,0 +1,87 @@ +""" +A Simulation is the environment of a Markov Game. +The Simulation is not responsible for properly checking / formatting the responses of LLM's. +This is the job of the `Agent` class. +Simulations expect clean actions, and are defined similarly to `gymnasium` environments, except that they are adapted for the Multi-agent setting. +""" + +from abc import ABC, abstractmethod +from typing import Any, Tuple + +from numpy.random import default_rng + +from mllm.markov_games.rollout_tree import SimulationStepLog + + +class Simulation(ABC): + @abstractmethod + def __init__(self, seed: int, *args, **kwargs): + self.seed = seed + self.rng = default_rng(self.seed) + + @abstractmethod + def step(self, actions: Any) -> Tuple[bool, SimulationStepLog]: + """ + Returns terminated, info + """ + raise NotImplementedError + + def get_obs(self): + """Returns all agent observations in dict + + Returns: + observations + """ + raise NotImplementedError + + def get_obs_agent(self, agent_id): + """Returns observation for agent_id""" + raise NotImplementedError + + def get_obs_size(self): + """Returns the shape of the observation""" + raise NotImplementedError + + def get_state(self): + raise NotImplementedError + + def get_state_size(self): + """Returns the shape of the state""" + raise NotImplementedError + + def get_avail_actions(self): + raise NotImplementedError + + def get_avail_agent_actions(self, agent_id): + """Returns the available actions for agent_id""" + raise NotImplementedError + + def get_total_actions(self): + """Returns the total number of actions an agent could ever take""" + # TODO: This is only suitable for a discrete 1 dimensional action space for each agent + raise NotImplementedError + + def get_safe_copy(self): + """ + Return copy of the agent object that is decorrelated from the original object. + """ + raise NotImplementedError + + def reset(self): + """Returns initial observations and states""" + raise NotImplementedError + + def render(self): + raise NotImplementedError + + def close(self): + raise NotImplementedError + + # def seed(self): + # raise NotImplementedError + + def save_replay(self): + raise NotImplementedError + + def get_simulation_info(self): + raise NotImplementedError diff --git a/src_code_for_reproducibility/markov_games/statistics_runner.py b/src_code_for_reproducibility/markov_games/statistics_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..eb30800f4aa4e5f919292023df77fa29e43d9c24 --- /dev/null +++ b/src_code_for_reproducibility/markov_games/statistics_runner.py @@ -0,0 +1,405 @@ +from __future__ import annotations + +import gc +import json +import pickle +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional + +from basic_render import find_iteration_folders + +from mllm.markov_games.rollout_tree import ( + RolloutTreeBranchNode, + RolloutTreeNode, + RolloutTreeRootNode, + SimulationStepLog, +) + + +def _iterate_main_nodes(root: RolloutTreeRootNode) -> Iterator[RolloutTreeNode]: + """ + Iterate the main path nodes without materializing full path lists. + """ + current = root.child + while current is not None: + if isinstance(current, RolloutTreeNode): + yield current + current = current.child + elif isinstance(current, RolloutTreeBranchNode): + # Follow only the main child on the main trajectory + current = current.main_child + else: + break + + +def iterate_main_simulation_logs( + root: RolloutTreeRootNode, +) -> Iterator[SimulationStepLog]: + for node in _iterate_main_nodes(root): + yield node.step_log.simulation_step_log + + +def stream_rollout_files(iteration_folder: Path) -> Iterator[Path]: + for p in iteration_folder.rglob("*.rt.pkl"): + if p.is_file(): + yield p + + +def load_root(path: Path) -> RolloutTreeRootNode: + with open(path, "rb") as f: + data = pickle.load(f) + return RolloutTreeRootNode.model_validate(data) + + +@dataclass +class StatRecord: + mgid: int + crn_id: Optional[int] + iteration: str + values: Dict[str, Any] + + +class StatComputer: + """ + Stateful stat computer that consumes SimulationStepLog instances + and produces final aggregated values for one rollout (mgid). + """ + + def update(self, sl: SimulationStepLog) -> None: # pragma: no cover - interface + raise NotImplementedError + + def finalize(self) -> Dict[str, Any]: # pragma: no cover - interface + raise NotImplementedError + + +def run_stats( + data_root: Path, + game_name: str, + make_computers: Callable[[], List[StatComputer]], + output_filename: Optional[str] = None, + output_format: str = "json", # "json" (dict of lists) or "jsonl" +) -> Path: + """ + Compute stats across all iteration_* folders under data_root. + Writes JSONL to data_root/statistics/. + """ + data_root = Path(data_root) + outdir = data_root / "statistics" + outdir.mkdir(parents=True, exist_ok=True) + # Choose extension by format + default_name = ( + f"{game_name}.stats.json" + if output_format == "json" + else f"{game_name}.stats.jsonl" + ) + outfile = outdir / ( + output_filename if output_filename is not None else default_name + ) + + # Rewrite file each run to keep it clean and small + if outfile.exists(): + outfile.unlink() + + iteration_folders = find_iteration_folders(str(data_root)) + + # If writing JSONL, stream directly; otherwise accumulate minimal records + if output_format == "jsonl": + with open(outfile, "w", encoding="utf-8") as w: + for iteration_folder in iteration_folders: + iteration_name = Path(iteration_folder).name + for pkl_path in stream_rollout_files(Path(iteration_folder)): + root = load_root(pkl_path) + + computers = make_computers() + for sl in iterate_main_simulation_logs(root): + for comp in computers: + try: + comp.update(sl) + except Exception: + continue + + values: Dict[str, Any] = {} + for comp in computers: + try: + values.update(comp.finalize()) + except Exception: + continue + + rec = { + "mgid": getattr(root, "id", None), + "crn_id": getattr(root, "crn_id", None), + "iteration": iteration_name, + "stats": values, + } + w.write(json.dumps(rec, ensure_ascii=False) + "\n") + + del root + del computers + gc.collect() + else: + # Aggregate to dict-of-lists for easier plotting + records: List[Dict[str, Any]] = [] + # Process in deterministic order + for iteration_folder in iteration_folders: + iteration_name = Path(iteration_folder).name + for pkl_path in stream_rollout_files(Path(iteration_folder)): + root = load_root(pkl_path) + + computers = make_computers() + for sl in iterate_main_simulation_logs(root): + for comp in computers: + try: + comp.update(sl) + except Exception: + continue + + values: Dict[str, Any] = {} + for comp in computers: + try: + values.update(comp.finalize()) + except Exception: + continue + + records.append( + { + "mgid": getattr(root, "id", None), + "crn_id": getattr(root, "crn_id", None), + "iteration": iteration_name, + "stats": values, + } + ) + + del root + del computers + gc.collect() + + # Build dict-of-lists with nested stats preserved + # Collect all stat keys and nested agent keys where needed + mgids: List[Any] = [] + crn_ids: List[Any] = [] + iterations_out: List[str] = [] + # stats_out is a nested structure mirroring keys but with lists + stats_out: Dict[str, Any] = {} + + # First pass to collect union of keys + stat_keys: set[str] = set() + nested_agent_keys: Dict[str, set[str]] = {} + for r in records: + stats = r.get("stats", {}) or {} + for k, v in stats.items(): + stat_keys.add(k) + if isinstance(v, dict): + nested = nested_agent_keys.setdefault(k, set()) + for ak in v.keys(): + nested.add(str(ak)) + + # Initialize structure + for k in stat_keys: + if k in nested_agent_keys: + stats_out[k] = {ak: [] for ak in sorted(nested_agent_keys[k])} + else: + stats_out[k] = [] + + # Fill lists + for r in records: + mgids.append(r.get("mgid")) + crn_ids.append(r.get("crn_id")) + iterations_out.append(r.get("iteration")) + stats = r.get("stats", {}) or {} + for k in stat_keys: + val = stats.get(k) + if isinstance(stats_out[k], dict): + # per-agent dict + agent_dict = val if isinstance(val, dict) else {} + for ak in stats_out[k].keys(): + stats_out[k][ak].append(agent_dict.get(ak)) + else: + stats_out[k].append(val) + + with open(outfile, "w", encoding="utf-8") as w: + json.dump( + { + "mgid": mgids, + "crn_id": crn_ids, + "iteration": iterations_out, + "stats": stats_out, + }, + w, + ensure_ascii=False, + ) + + return outfile + + +def run_stats_functional( + data_root: Path, + game_name: str, + metrics: Dict[str, Callable[[SimulationStepLog], Optional[Dict[str, float]]]], + output_filename: Optional[str] = None, + output_format: str = "json", +) -> Path: + """ + Functional variant where metrics is a dict of name -> f(SimulationStepLog) -> {agent_id: value}. + Aggregates per rollout by averaging over steps where a metric produced a value. + Writes a single consolidated file in data_root/statistics/. + """ + data_root = Path(data_root) + outdir = data_root / "statistics" + outdir.mkdir(parents=True, exist_ok=True) + default_name = ( + f"{game_name}.stats.json" + if output_format == "json" + else f"{game_name}.stats.jsonl" + ) + outfile = outdir / ( + output_filename if output_filename is not None else default_name + ) + + if outfile.exists(): + outfile.unlink() + + iteration_folders = find_iteration_folders(str(data_root)) + + def finalize_rollout( + agg: Dict[str, Dict[str, List[float]]] + ) -> Dict[str, Dict[str, float]]: + # avg per metric per agent + result: Dict[str, Dict[str, float]] = {} + for mname, agent_values in agg.items(): + result[mname] = {} + for aid, vals in agent_values.items(): + if not vals: + result[mname][aid] = None # keep alignment; could be None + else: + result[mname][aid] = sum(vals) / len(vals) + return result + + if output_format == "jsonl": + with open(outfile, "w", encoding="utf-8") as w: + for iteration_folder in iteration_folders: + iteration_name = Path(iteration_folder).name + for pkl_path in stream_rollout_files(Path(iteration_folder)): + root = load_root(pkl_path) + + # aggregator structure: metric -> agent_id -> list of values + agg: Dict[str, Dict[str, List[float]]] = { + m: {} for m in metrics.keys() + } + + for sl in iterate_main_simulation_logs(root): + for mname, fn in metrics.items(): + try: + vals = fn(sl) + except Exception: + vals = None + if not vals: + continue + for aid, v in vals.items(): + if v is None: + continue + lst = agg[mname].setdefault(str(aid), []) + try: + lst.append(float(v)) + except Exception: + continue + + values = finalize_rollout(agg) + rec = { + "mgid": getattr(root, "id", None), + "crn_id": getattr(root, "crn_id", None), + "iteration": iteration_name, + "stats": values, + } + w.write(json.dumps(rec, ensure_ascii=False) + "\n") + + del root + gc.collect() + else: + records: List[Dict[str, Any]] = [] + for iteration_folder in iteration_folders: + iteration_name = Path(iteration_folder).name + for pkl_path in stream_rollout_files(Path(iteration_folder)): + root = load_root(pkl_path) + + agg: Dict[str, Dict[str, List[float]]] = {m: {} for m in metrics.keys()} + for sl in iterate_main_simulation_logs(root): + for mname, fn in metrics.items(): + try: + vals = fn(sl) + except Exception: + vals = None + if not vals: + continue + for aid, v in vals.items(): + if v is None: + continue + lst = agg[mname].setdefault(str(aid), []) + try: + lst.append(float(v)) + except Exception: + continue + + values = finalize_rollout(agg) + records.append( + { + "mgid": getattr(root, "id", None), + "crn_id": getattr(root, "crn_id", None), + "iteration": iteration_name, + "stats": values, + } + ) + + del root + gc.collect() + + # Build dict-of-lists output + mgids: List[Any] = [] + crn_ids: List[Any] = [] + iterations_out: List[str] = [] + stats_out: Dict[str, Any] = {} + + stat_keys: set[str] = set() + nested_agent_keys: Dict[str, set[str]] = {} + for r in records: + stats = r.get("stats", {}) or {} + for k, v in stats.items(): + stat_keys.add(k) + if isinstance(v, dict): + nested = nested_agent_keys.setdefault(k, set()) + for ak in v.keys(): + nested.add(str(ak)) + + for k in stat_keys: + if k in nested_agent_keys: + stats_out[k] = {ak: [] for ak in sorted(nested_agent_keys[k])} + else: + stats_out[k] = [] + + for r in records: + mgids.append(r.get("mgid")) + crn_ids.append(r.get("crn_id")) + iterations_out.append(r.get("iteration")) + stats = r.get("stats", {}) or {} + for k in stat_keys: + val = stats.get(k) + if isinstance(stats_out[k], dict): + agent_dict = val if isinstance(val, dict) else {} + for ak in stats_out[k].keys(): + stats_out[k][ak].append(agent_dict.get(ak)) + else: + stats_out[k].append(val) + + with open(outfile, "w", encoding="utf-8") as w: + json.dump( + { + "mgid": mgids, + "crn_id": crn_ids, + "iteration": iterations_out, + "stats": stats_out, + }, + w, + ensure_ascii=False, + ) + + return outfile diff --git a/src_code_for_reproducibility/markov_games/vine_ppo.py b/src_code_for_reproducibility/markov_games/vine_ppo.py new file mode 100644 index 0000000000000000000000000000000000000000..f3cd2c89133fa00b0d1d0fa260688efa642d4ded --- /dev/null +++ b/src_code_for_reproducibility/markov_games/vine_ppo.py @@ -0,0 +1,10 @@ +from anytree import Node, RenderTree +from anytree.exporter import DotExporter +import os.path +import asyncio +from mllm.markov_games.markov_game import MarkovGame + +async def VinePPORunner( + markov_game: MarkovGame, + **kwargs): + pass diff --git a/src_code_for_reproducibility/models/__init__.py b/src_code_for_reproducibility/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src_code_for_reproducibility/models/adapter_training_wrapper.py b/src_code_for_reproducibility/models/adapter_training_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..ccb41c9fb2c314db09a3e8a370ffa50c3648a198 --- /dev/null +++ b/src_code_for_reproducibility/models/adapter_training_wrapper.py @@ -0,0 +1,98 @@ +import torch +import torch.nn as nn +import logging +from typing import Union +from peft import ( + LoraConfig, + get_peft_model, +) + +logger = logging.getLogger(__name__) + + +class AdapterWrapper(nn.Module): + """ + A thin faΓ§ade that + β€’ keeps a reference to a *shared* PEFT-wrapped model, + β€’ ensures `set_adapter(adapter)` is called on every forward, + β€’ exposes only the parameters that should be trained for that adapter + (plus whatever extra modules you name). + """ + def __init__( + self, + shared_llm: nn.Module, + adapter_id: str, + lora_config: dict, + path: Union[str, None] = None, + ): + super().__init__() + self.shared_llm = shared_llm + self.adapter_id = adapter_id + lora_config = LoraConfig(**lora_config) + # this modifies the shared llm in place, adding a lora adapter inside + self.shared_llm = get_peft_model( + model=shared_llm, + peft_config=lora_config, + adapter_name=adapter_id, + ) + self.shared_llm.train() + # Load external adapter weights if provided + loaded_from: str | None = None + if path: + try: + # Supports both local filesystem paths and HF Hub repo IDs + self.shared_llm.load_adapter( + is_trainable=True, + model_id=path, + adapter_name=adapter_id, + ) + loaded_from = path + except Exception as exc: # noqa: BLE001 - want to log any load failure context + logger.warning( + f"Adapter '{adapter_id}': failed to load from '{path}': {exc}" + ) + + if loaded_from: + logger.info( + f"Adapter '{adapter_id}': loaded initial weights from '{loaded_from}'." + ) + else: + logger.info( + f"Adapter '{adapter_id}': initialized with fresh weights (no initial weights found)." + ) + + def parameters(self, recurse: bool = True): + """ + "recurse" is just for pytorch compatibility + """ + self.shared_llm.set_adapter(self.adapter_id) + params = [p for p in self.shared_llm.parameters() if p.requires_grad] + + return params + + def get_base_model_logits(self, contexts): + """ + Run the base model (without adapter) in inference mode, without tracking gradients. + This is useful to get reference logits for KL-divergence computation. + """ + with torch.no_grad(): + with self.shared_llm.disable_adapter(): + return self.shared_llm(input_ids=contexts)[0] + + def forward(self, *args, **kwargs): + self.shared_llm.set_adapter(self.adapter_id) + return self.shared_llm(*args, **kwargs) + + def save_pretrained(self, save_path): + self.shared_llm.save_pretrained(save_path) + + def gradient_checkpointing_enable(self, *args, **kwargs): + self.shared_llm.gradient_checkpointing_enable(*args, **kwargs) + + @property + def dtype(self): + return self.shared_llm.dtype + + @property + def device(self): + return self.shared_llm.device diff --git a/src_code_for_reproducibility/models/human_policy.py b/src_code_for_reproducibility/models/human_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..00076a9e592cb29f8e6f33d852fc51be3e47ed13 --- /dev/null +++ b/src_code_for_reproducibility/models/human_policy.py @@ -0,0 +1,255 @@ +import asyncio +import os +import re +import shutil +import sys +from typing import Callable, Dict, List, Optional + +from mllm.markov_games.rollout_tree import ChatTurn + +try: + import rstr # For generating example strings from regex +except Exception: # pragma: no cover + rstr = None + + +def _clear_terminal() -> None: + """ + Clear the terminal screen in a cross-platform manner. + """ + if sys.stdout.isatty(): + os.system("cls" if os.name == "nt" else "clear") + + +def _terminal_width(default: int = 100) -> int: + try: + return shutil.get_terminal_size().columns + except Exception: + return default + + +def _horizontal_rule(char: str = "─") -> str: + width = max(20, _terminal_width() - 2) + return char * width + + +class _Style: + # ANSI colors (bright, readable) + RESET = "\033[0m" + BOLD = "\033[1m" + DIM = "\033[2m" + # Foreground colors + FG_BLUE = "\033[94m" # user/system headers + FG_GREEN = "\033[92m" # human response header + FG_YELLOW = "\033[93m" # notices + FG_RED = "\033[91m" # errors + FG_MAGENTA = "\033[95m" # regex + FG_CYAN = "\033[96m" # tips + + +def _render_chat(state) -> str: + """ + Render prior messages in a compact, readable terminal format. + + Expected message dict keys: {"role": str, "content": str, ...} + """ + lines: List[str] = [] + lines.append(_horizontal_rule()) + lines.append(f"{_Style.FG_BLUE}{_Style.BOLD} Conversation so far {_Style.RESET}") + lines.append(_horizontal_rule()) + for chat in state: + role = chat.role + content = str(chat.content).strip() + # Map roles to display names and colors/emojis + if role == "assistant": + header = f"{_Style.FG_GREEN}{_Style.BOLD}HUMAN--πŸ§‘β€πŸ’»{_Style.RESET}" + elif role == "user": + header = f"{_Style.FG_BLUE}{_Style.BOLD}USER--βš™οΈ{_Style.RESET}" + else: + header = f"[{_Style.DIM}{role.upper()}{_Style.RESET}]" + lines.append(header) + # Indent content for readability + for line in content.splitlines() or [""]: + lines.append(f" {line}") + lines.append("") + lines.append(_horizontal_rule()) + return "\n".join(lines) + + +async def _async_input(prompt_text: str) -> str: + """Non-blocking input using a background thread.""" + return await asyncio.to_thread(input, prompt_text) + + +def _short_regex_example(regex: str, max_len: int = 30) -> Optional[str]: + """ + Try to produce a short example string that matches the regex. + We attempt multiple times and pick the first <= max_len. + """ + if rstr is None: + return None + try: + for _ in range(20): + candidate = rstr.xeger(regex) + if len(candidate) <= max_len: + return candidate + # Fallback to truncation (may break match, so don't return) + return None + except Exception: + return None + + +def _detect_input_type(regex: str | None) -> tuple[str, str, str]: + """ + Detect what type of input is expected based on the regex pattern. + Returns (input_type, start_tag, end_tag) + """ + if regex is None: + return "text", "", "" + + if "message_start" in regex and "message_end" in regex: + return "message", "<>", "<>" + elif "proposal_start" in regex and "proposal_end" in regex: + return "proposal", "<>", "<>" + else: + return "text", "", "" + + +async def human_policy(state, agent_id, regex: str | None = None) -> str: + """ + Async human-in-the-loop policy. + + - Displays prior conversation context in the terminal. + - Prompts the user for a response. + - If a regex is provided, validates and re-prompts until it matches. + - Automatically adds formatting tags based on expected input type. + + Args: + prompt: Chat history as a list of {role, content} dicts. + regex: Optional fullmatch validation pattern. + + Returns: + The user's validated response string. + """ + # Detect input type and formatting + input_type, start_tag, end_tag = _detect_input_type(regex) + + while True: + _clear_terminal() + print(_render_chat(state)) + + if regex: + example = _short_regex_example(regex, max_len=30) + print( + f"{_Style.FG_MAGENTA}{_Style.BOLD}Expected format (regex fullmatch):{_Style.RESET}" + ) + print(f" {_Style.FG_MAGENTA}{regex}{_Style.RESET}") + if example: + print( + f"{_Style.FG_CYAN}Example (random, <=30 chars):{_Style.RESET} {example}" + ) + print(_horizontal_rule(".")) + + # Custom prompt based on input type + if input_type == "message": + print( + f"{_Style.FG_YELLOW}Type your message content (formatting will be added automatically):{_Style.RESET}" + ) + elif input_type == "proposal": + print( + f"{_Style.FG_YELLOW}Type your proposal (number only, formatting will be added automatically):{_Style.RESET}" + ) + else: + print( + f"{_Style.FG_YELLOW}Type your response and press Enter.{_Style.RESET}" + ) + + print( + f"{_Style.DIM}Commands: /help to view commands, /refresh to re-render, /quit to abort{_Style.RESET}" + ) + else: + print( + f"{_Style.FG_YELLOW}Type your response and press Enter.{_Style.RESET} {_Style.DIM}(/help for commands){_Style.RESET}" + ) + + user_in = (await _async_input("> ")).rstrip("\n") + + # Commands + if user_in.strip().lower() in {"/help", "/h"}: + print(f"\n{_Style.FG_CYAN}{_Style.BOLD}Available commands:{_Style.RESET}") + print( + f" {_Style.FG_CYAN}/help{_Style.RESET} or {_Style.FG_CYAN}/h{_Style.RESET} Show this help" + ) + print( + f" {_Style.FG_CYAN}/refresh{_Style.RESET} or {_Style.FG_CYAN}/r{_Style.RESET} Re-render the conversation and prompt" + ) + print( + f" {_Style.FG_CYAN}/quit{_Style.RESET} or {_Style.FG_CYAN}/q{_Style.RESET} Abort the run (raises KeyboardInterrupt)" + ) + await asyncio.sleep(1.0) + continue + if user_in.strip().lower() in {"/refresh", "/r"}: + continue + if user_in.strip().lower() in {"/quit", "/q"}: + raise KeyboardInterrupt("Human aborted run from human_policy") + + # Add formatting tags if needed + if start_tag and end_tag: + formatted_input = f"{start_tag}{user_in}{end_tag}" + else: + formatted_input = user_in + + if regex is None: + return ChatTurn( + role="assistant", agent_id=agent_id, content=formatted_input + ) + + # Validate against regex (fullmatch) + try: + pattern = re.compile(regex) + except re.error as e: + # If regex is invalid, fall back to accepting any input + print( + f"{_Style.FG_RED}Warning:{_Style.RESET} Provided regex is invalid: {e}. Accepting input without validation." + ) + await asyncio.sleep(0.5) + return ChatTurn( + role="assistant", agent_id=agent_id, content=formatted_input + ) + + if pattern.fullmatch(formatted_input): + return ChatTurn( + role="assistant", agent_id=agent_id, content=formatted_input + ) + + # Show validation error and re-prompt + print("") + print( + f"{_Style.FG_RED}{_Style.BOLD}Input did not match the required format.{_Style.RESET} Please try again." + ) + + if input_type == "message": + print( + f"You entered: {_Style.FG_CYAN}{start_tag}{user_in}{end_tag}{_Style.RESET}" + ) + print(f"Just type the message content without tags.") + elif input_type == "proposal": + print( + f"You entered: {_Style.FG_CYAN}{start_tag}{user_in}{end_tag}{_Style.RESET}" + ) + print(f"Just type the number without tags.") + else: + print(f"Expected (regex):") + print(f" {_Style.FG_MAGENTA}{regex}{_Style.RESET}") + + print(_horizontal_rule(".")) + print(f"{_Style.FG_YELLOW}Press Enter to retry...{_Style.RESET}") + await _async_input("") + + +def get_human_policies() -> Dict[str, Callable[[List[Dict]], str]]: + """ + Expose the human policy in the same map shape used elsewhere. + """ + # Type hint says Callable[[List[Dict]], str] but we intentionally return the async callable. + return {"human_policy": human_policy} # type: ignore[return-value] diff --git a/src_code_for_reproducibility/models/inference_backend.py b/src_code_for_reproducibility/models/inference_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..2fd31f0f51e00dbcbf1db954e8eab42550e05d4e --- /dev/null +++ b/src_code_for_reproducibility/models/inference_backend.py @@ -0,0 +1,39 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Optional + + +@dataclass +class LLMInferenceOutput: + content: str + reasoning_content: str | None = None + log_probs: list[float] | None = None + out_token_ids: list[int] | None = None + + +class LLMInferenceBackend(ABC): + @abstractmethod + def __init__(self, **kwargs): + ... + + @abstractmethod + def prepare_adapter( + self, adapter_id: str, weights_got_updated: bool = False + ) -> None: + """Ensure adapter is ready/loaded for next generation call.""" + + @abstractmethod + async def generate(self, prompt: list[dict], regex: Optional[str] = None) -> str: + ... + + @abstractmethod + def toggle_training_mode(self) -> None: + ... + + @abstractmethod + def toggle_eval_mode(self) -> None: + ... + + @abstractmethod + def shutdown(self) -> None: + ... diff --git a/src_code_for_reproducibility/models/inference_backend_dummy.py b/src_code_for_reproducibility/models/inference_backend_dummy.py new file mode 100644 index 0000000000000000000000000000000000000000..ac7d4bd5566220377913e6af012545df492acad9 --- /dev/null +++ b/src_code_for_reproducibility/models/inference_backend_dummy.py @@ -0,0 +1,54 @@ +import asyncio +from typing import Optional + +import rstr +from transformers import AutoTokenizer + +from mllm.models.inference_backend import LLMInferenceBackend, LLMInferenceOutput +from mllm.utils.short_id_gen import generate_short_id + + +class DummyInferenceBackend(LLMInferenceBackend): + def __init__( + self, + *args, + **kwargs, + ): + pass + + def prepare_adapter( + self, + adapter_id: Optional[str], + weights_got_updated: bool, + adapter_path: Optional[str] = None, + ) -> None: + pass + + async def toggle_training_mode(self) -> None: + await asyncio.sleep(0) + pass + + async def toggle_eval_mode(self) -> None: + await asyncio.sleep(0) + pass + + def shutdown(self) -> None: + pass + + async def generate( + self, + prompt_text: str, + regex: Optional[str] = None, + extract_thinking: bool = False, + ) -> LLMInferenceOutput: + if regex: + # Create random string that respects the regex + return LLMInferenceOutput( + content=rstr.xeger(regex), + reasoning_content="I don't think, I am a dummy backend.", + ) + else: + return LLMInferenceOutput( + content="I am a dummy backend without a regex.", + reasoning_content="I don't think, I am a dummy backend.", + ) diff --git a/src_code_for_reproducibility/models/inference_backend_sglang.py b/src_code_for_reproducibility/models/inference_backend_sglang.py new file mode 100644 index 0000000000000000000000000000000000000000..d06b9d4bed25bb2b7fe117422b67b40f0c6e4509 --- /dev/null +++ b/src_code_for_reproducibility/models/inference_backend_sglang.py @@ -0,0 +1,86 @@ +# new_backend_sglang_offline.py +from __future__ import annotations + +import asyncio +from typing import Any, Optional + +# import sglang as sgl + +from mllm.models.inference_backend import LLMInferenceBackend + + +class SGLangOfflineBackend(LLMInferenceBackend): + def __init__( + self, + model_name: str, + tokenizer, # unused but kept for parity + adapter_paths: dict[str, str], + device: str = "cuda", + max_model_len: Optional[int] = None, + enable_lora: bool = True, + lora_target_modules: Optional[list[str] | str] = None, + max_loras_per_batch: int = 8, + engine_kwargs: dict[str, Any] = None, + ): + self.model_name = model_name + self.adapter_paths = adapter_paths + self.current_adapter: Optional[str] = None + engine_kwargs = dict(engine_kwargs or {}) + # Map server-style LoRA flags to offline engine ctor + if enable_lora and adapter_paths: + engine_kwargs.setdefault("enable_lora", True) + # The offline Engine mirrors server args; pass a mapping name->path + engine_kwargs.setdefault("lora_paths", adapter_paths) + if lora_target_modules is not None: + engine_kwargs.setdefault("lora_target_modules", lora_target_modules) + engine_kwargs.setdefault("max_loras_per_batch", max_loras_per_batch) + + if max_model_len is not None: + engine_kwargs.setdefault("context_length", max_model_len) + + # Launch in-process engine (no HTTP server) + self.llm = sgl.Engine(model_path=model_name, **engine_kwargs) # async-ready + # SGLang supports: generate(), async_generate(), and async streaming helpers. :contentReference[oaicite:2]{index=2} + + def is_ready(self) -> bool: + return True + + def toggle_training_mode(self) -> None: + # No explicit KV release API offline; typically you pause usage here. + pass + + def toggle_eval_mode(self) -> None: + pass + + def shutdown(self) -> None: + # Engine cleans up on GC; explicit close not required. + pass + + def prepare_adapter(self, adapter_id: Optional[str]) -> None: + # With offline Engine, when LoRA is enabled at init, + # you select adapter per request via the input batch mapping. + self.current_adapter = adapter_id + + async def generate( + self, prompt_text: str, sampling_params: dict, adapter_id: Optional[str] + ) -> str: + # Non-streaming async (batch of 1). For batched prompts, pass a list. + params = { + "temperature": sampling_params.get("temperature", 1.0), + "top_p": sampling_params.get("top_p", 1.0), + "max_new_tokens": sampling_params.get("max_new_tokens", 128), + } + if (tk := sampling_params.get("top_k", -1)) and tk > 0: + params["top_k"] = tk + if (mn := sampling_params.get("min_new_tokens")) is not None: + params["min_new_tokens"] = mn + if (fp := sampling_params.get("frequency_penalty")) is not None: + params["frequency_penalty"] = fp + + # If using multi-LoRA, SGLang lets you provide adapter names aligned to each input. + prompts = [prompt_text] + adapters = [adapter_id] if adapter_id else None # or omit for base + outs = await self.llm.async_generate( + prompts, params, adapters + ) # :contentReference[oaicite:3]{index=3} + return outs[0]["text"] diff --git a/src_code_for_reproducibility/models/inference_backend_sglang_local_server.py b/src_code_for_reproducibility/models/inference_backend_sglang_local_server.py new file mode 100644 index 0000000000000000000000000000000000000000..c29f4d01a0bc0a6a0435461a02cd67074771abb0 --- /dev/null +++ b/src_code_for_reproducibility/models/inference_backend_sglang_local_server.py @@ -0,0 +1,127 @@ +import os + +import httpx +import requests +from sglang.utils import launch_server_cmd, wait_for_server + +from mllm.models.inference_backend import LLMInferenceBackend + + +class HttpSGLangBackend(LLMInferenceBackend): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.port = None + self.proc = None + self.urls = {} + # track sglang adapter ids separately from your logical ids + self.sglang_names = {aid: aid for aid in self.adapter_paths.keys()} + self.needs_loading = {aid: True for aid in self.adapter_paths.keys()} + + # defaults you already used: + self.mem_fraction = kwargs.get("mem_fraction_static", 0.6) + self.dtype = kwargs.get("dtype", "bfloat16") + self.extra_cli = kwargs.get("extra_cli", "") + self.disable_radix_cache = kwargs.get("disable_radix_cache", True) + + def launch(self) -> None: + # find local hf cache path for server + from transformers.utils import cached_file + + local_llm_path = os.path.split(cached_file(self.model_name, "config.json"))[0] + + lora_str = "" + if self.adapter_paths: + lora_str = "--lora-paths " + " ".join( + f"{aid}={path}" for aid, path in self.adapter_paths.items() + ) + + cmd = f""" + python3 -m sglang.launch_server --model-path {local_llm_path} \ + --host 0.0.0.0 {lora_str} \ + {'--disable-radix-cache' if self.disable_radix_cache else ''} \ + --mem-fraction-static {self.mem_fraction} --dtype {self.dtype} {self.extra_cli} + """ + self.proc, self.port = launch_server_cmd(cmd) + wait_for_server(f"http://localhost:{self.port}") + base = f"http://localhost:{self.port}" + self.urls = dict( + generate=f"{base}/generate", + release=f"{base}/release_memory_occupation", + resume=f"{base}/resume_memory_occupation", + load_lora=f"{base}/load_lora_adapter", + unload_lora=f"{base}/unload_lora_adapter", + ) + + def is_ready(self) -> bool: + try: + requests.get(self.urls["generate"], timeout=2) + return True + except Exception: + return False + + def prepare_adapter(self, adapter_id: str) -> None: + if adapter_id is None: + return + if self.needs_loading.get(adapter_id, False): + # unload old name if present + try: + requests.post( + self.urls["unload_lora"], + json={"lora_name": self.sglang_names[adapter_id]}, + timeout=10, + ) + except Exception: + pass + new_name = self._short_id() + self.sglang_names[adapter_id] = new_name + requests.post( + self.urls["load_lora"], + json={ + "lora_name": new_name, + "lora_path": self.adapter_paths[adapter_id], + }, + ).raise_for_status() + self.needs_loading[adapter_id] = False + + async def generate( + self, prompt_text: str, sampling_params: dict, adapter_id: str | None + ) -> str: + lora_name = self.sglang_names.get(adapter_id) if adapter_id else None + payload = { + "text": [prompt_text], + "sampling_params": sampling_params, + } + if lora_name: + payload["lora_path"] = [lora_name] + + timeout = httpx.Timeout(3600.0, connect=3600.0) + async with httpx.AsyncClient(timeout=timeout) as client: + resp = await client.post(self.urls["generate"], json=payload) + resp.raise_for_status() + return resp.json()[0]["text"] + + def toggle_training_mode(self) -> None: + # free KV space while training adapters + requests.post( + self.urls["release"], json={"tags": ["kv_cache"]} + ).raise_for_status() + + def toggle_eval_mode(self) -> None: + # re-allocate KV space + try: + requests.post( + self.urls["resume"], json={"tags": ["kv_cache"]} + ).raise_for_status() + except Exception: + pass + + def shutdown(self) -> None: + from sglang.utils import terminate_process + + if self.proc: + terminate_process(self.proc) + + def _short_id(self) -> str: + import uuid + + return str(uuid.uuid4().int)[:8] diff --git a/src_code_for_reproducibility/models/inference_backend_vllm.py b/src_code_for_reproducibility/models/inference_backend_vllm.py new file mode 100644 index 0000000000000000000000000000000000000000..fed694711057ba831e57a75b7b53d019026d835a --- /dev/null +++ b/src_code_for_reproducibility/models/inference_backend_vllm.py @@ -0,0 +1,117 @@ +import asyncio +import re +from typing import Optional + +import torch +from transformers import AutoTokenizer +from vllm import AsyncEngineArgs, AsyncLLMEngine, SamplingParams +from vllm.inputs import TokensPrompt +from vllm.lora.request import LoRARequest +from vllm.sampling_params import GuidedDecodingParams, RequestOutputKind + +from mllm.models.inference_backend import LLMInferenceBackend, LLMInferenceOutput +from mllm.utils.short_id_gen import generate_short_id + + +class VLLMAsyncBackend(LLMInferenceBackend): + def __init__( + self, + model_name: str, + tokenizer: AutoTokenizer, + # adapter_paths: dict[str, str], + engine_init_kwargs: dict = {}, + sampling_params: dict = {}, + ): + self.model_name = model_name + # self.adapter_paths = adapter_paths or {} + # self.current_adapter = None + # self.vllm_adapter_ids = { + # adapter_id: generate_short_id() for adapter_id in adapter_paths.keys() + # } + self.vllm_adapter_ids = {} + ea = dict(model=model_name, **engine_init_kwargs) + # ea["enable_lora"] = True + # ea["max_loras"] = len(self.vllm_adapter_ids) + # ea["enable_sleep_mode"] = True + self.engine = AsyncLLMEngine.from_engine_args(AsyncEngineArgs(**ea)) + + self.sampling_params = sampling_params + + def prepare_adapter( + self, + adapter_id: Optional[str], + adapter_path: Optional[str], + weights_got_updated: bool, + ) -> None: + # self.current_adapter = adapter_id + if weights_got_updated: + self.vllm_adapter_ids[adapter_id] = generate_short_id() + self.current_lora_request = LoRARequest( + adapter_id, + self.vllm_adapter_ids[adapter_id], + adapter_path, + ) + + async def toggle_training_mode(self) -> None: + await self.engine.sleep(level=1) + + async def toggle_eval_mode(self) -> None: + await self.engine.wake_up() + + def shutdown(self) -> None: + # No explicit close call; engine stops when process exits. + pass + + async def generate( + self, + input_token_ids: list[int], + regex: Optional[str] = None, + extract_thinking: bool = False, + ) -> LLMInferenceOutput: + # Build SamplingParams correctly + guided = GuidedDecodingParams(regex=regex) if regex else None + sp = SamplingParams( + **self.sampling_params, + guided_decoding=guided, + output_kind=RequestOutputKind.FINAL_ONLY, + ) + + prompt = TokensPrompt(prompt_token_ids=input_token_ids) + request_id = f"req-{asyncio.get_running_loop().time()}" + result_generator = self.engine.generate( + prompt, + sp, # SamplingParams(...) + request_id, + lora_request=self.current_lora_request, + ) + + async for out in result_generator: # with FINAL_ONLY this runs once + res = out + + raw_text = res.outputs[0].text + out_token_ids = res.outputs[0].token_ids + log_probs = [ + logprob_dict[token_id].logprob + for token_id, logprob_dict in zip(out_token_ids, res.outputs[0].logprobs) + ] + log_probs = torch.tensor(log_probs) + out_token_ids = torch.tensor(out_token_ids, dtype=torch.long) + # for out_token_id, logprob_dict in zip(out_token_ids, res.outputs[0].logprobs): + # if logprob_dict[out_token_id].logprob < -1: + # print(f"High negative logprob {logprob_dict[out_token_id].logprob} for {logprob_dict}") + content = raw_text + reasoning_content = None + + if extract_thinking: + m = re.match( + r"^\n\n([\s\S]*?)\n\n(.*)$", raw_text, flags=re.DOTALL + ) + if m: + reasoning_content = m.group(1) + content = m.group(2) + return LLMInferenceOutput( + content=content, + reasoning_content=reasoning_content, + log_probs=log_probs, + out_token_ids=out_token_ids, + ) diff --git a/src_code_for_reproducibility/models/inference_backend_vllm_local_server.py b/src_code_for_reproducibility/models/inference_backend_vllm_local_server.py new file mode 100644 index 0000000000000000000000000000000000000000..815ba3eb2cc9bb1d664b0065b240cd0273f66474 --- /dev/null +++ b/src_code_for_reproducibility/models/inference_backend_vllm_local_server.py @@ -0,0 +1,160 @@ +import json +import os +import subprocess +import time + +import httpx +import requests + +from mllm.models.inference_backend import LLMInferenceBackend + + +class HttpVLLMBackend(LLMInferenceBackend): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.port = kwargs.get("port", 8000) + self.host = kwargs.get("host", "0.0.0.0") + self.proc = None + self.base_url = f"http://{self.host}:{self.port}" + # vLLM memory safety knobs + self.gpu_mem_util = kwargs.get("gpu_memory_utilization", 0.9) + self.max_model_len = kwargs.get("max_model_len", None) + self.max_num_seqs = kwargs.get("max_num_seqs", None) + self.max_batched_tokens = kwargs.get("max_num_batched_tokens", None) + self.dtype = kwargs.get("dtype", "bfloat16") + self.trust_remote_code = kwargs.get("trust_remote_code", False) + # LoRA strategy: "preload" (CLI) or "runtime" (endpoints) depending on your vLLM build + self.lora_mode = kwargs.get( + "lora_mode", "preload" + ) # "runtime" supported in newer builds + self.runtime_lora_enabled = self.lora_mode == "runtime" + + # If preloading: build CLI args (adapter name -> path) + self._preload_lora_args = [] + if self.adapter_paths and self.lora_mode == "preload": + # vLLM supports multiple LoRA modules via CLI in recent versions + # Example flag shapes can vary; adapt as needed for your version: + # --lora-modules adapter_id=path + for aid, pth in self.adapter_paths.items(): + self._preload_lora_args += ["--lora-modules", f"{aid}={pth}"] + + def launch(self): + # Build vLLM serve command + cmd = [ + "python3", + "-m", + "vllm.entrypoints.openai.api_server", + "--model", + self.model_name, + "--host", + self.host, + "--port", + str(self.port), + "--dtype", + self.dtype, + "--gpu-memory-utilization", + str(self.gpu_mem_util), + ] + if self.trust_remote_code: + cmd += ["--trust-remote-code"] + if self.max_model_len: + cmd += ["--max-model-len", str(self.max_model_len)] + if self.max_num_seqs: + cmd += ["--max-num-seqs", str(self.max_num_seqs)] + if self.max_batched_tokens: + cmd += ["--max-num-batched-tokens", str(self.max_batched_tokens)] + cmd += self._preload_lora_args + + self.proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True + ) + self._wait_ready() + + def _wait_ready(self, timeout=120): + url = f"{self.base_url}/v1/models" + t0 = time.time() + while time.time() - t0 < timeout: + try: + r = requests.get(url, timeout=2) + if r.status_code == 200: + return + except Exception: + pass + time.sleep(1) + raise RuntimeError("vLLM server did not become ready in time") + + def is_ready(self) -> bool: + try: + return ( + requests.get(f"{self.base_url}/v1/models", timeout=2).status_code == 200 + ) + except Exception: + return False + + def prepare_adapter(self, adapter_id: str) -> None: + if not adapter_id or not self.runtime_lora_enabled: + return + # Newer vLLM builds expose runtime LoRA endpoints. If yours differs, + # adjust the path/body here and keep the interface stable. + try: + requests.post( + f"{self.base_url}/v1/load_lora_adapter", + json={ + "adapter_name": adapter_id, + "adapter_path": self.adapter_paths[adapter_id], + }, + timeout=10, + ).raise_for_status() + except Exception as e: + # If already loaded or endpoint not present, swallow or log + pass + + async def generate( + self, prompt_text: str, sampling_params: dict, adapter_id: str | None + ) -> str: + # Map your sampling params to OpenAI schema + body = { + "model": self.model_name, + "messages": [{"role": "user", "content": prompt_text}], + "temperature": sampling_params.get("temperature", 1.0), + "top_p": sampling_params.get("top_p", 1.0), + "max_tokens": sampling_params.get("max_new_tokens", 128), + } + # Optional knobs: + if sampling_params.get("top_k", -1) and sampling_params["top_k"] > 0: + # vLLM accepts top_k via extra params; put under "extra_body" + body.setdefault("extra_body", {})["top_k"] = sampling_params["top_k"] + if sampling_params.get("min_new_tokens", None) is not None: + body.setdefault("extra_body", {})["min_tokens"] = sampling_params[ + "min_new_tokens" + ] + if sampling_params.get("frequency_penalty", None) is not None: + body["frequency_penalty"] = sampling_params["frequency_penalty"] + + # Select LoRA adapter + if adapter_id: + if self.runtime_lora_enabled: + body.setdefault("extra_body", {})["lora_adapter"] = adapter_id + else: + # when preloaded via CLI, most builds select by name via "adapter_name"/"lora_adapter" + body.setdefault("extra_body", {})["lora_adapter"] = adapter_id + + url = f"{self.base_url}/v1/chat/completions" + timeout = httpx.Timeout(3600.0, connect=3600.0) + async with httpx.AsyncClient(timeout=timeout) as client: + resp = await client.post(url, json=body) + resp.raise_for_status() + data = resp.json() + return data["choices"][0]["message"]["content"] + + def toggle_training_mode(self) -> None: + # vLLM doesn’t expose an explicit KV β€œrelease” toggle via API. + # Strategy: keep inference server idle during training, or run training in a separate process. + pass + + def toggle_eval_mode(self) -> None: + pass + + def shutdown(self) -> None: + if self.proc: + self.proc.terminate() diff --git a/src_code_for_reproducibility/models/large_language_model_api.py b/src_code_for_reproducibility/models/large_language_model_api.py new file mode 100644 index 0000000000000000000000000000000000000000..1afedf4e19f6b00c40c0127b18843eab03a89648 --- /dev/null +++ b/src_code_for_reproducibility/models/large_language_model_api.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import asyncio +import copy +import os +import random +import re +from typing import Any, Callable, Dict, List, Optional, Sequence + +import backoff +from openai import AsyncOpenAI, OpenAIError + +from mllm.markov_games.rollout_tree import ChatTurn +from mllm.models.inference_backend import LLMInferenceOutput + +# TODO: Get this automatically from OpenAI +reasoning_models = [ + "gpt-5-nano", + "gpt-5-mini", + "gpt-5", + "o1-mini", + "o1", + "o1-pro", + "o3-mini", + "o3", + "o3-pro", + "o4-mini", + "o4", + "o4-pro", +] + + +class LargeLanguageModelOpenAI: + """Tiny async wrapper for OpenAI Chat Completions.""" + + def __init__( + self, + llm_id: str = "", + model: str = "gpt-4.1-mini", + api_key: Optional[str] = None, + base_url: Optional[str] = None, + timeout_s: float = 300.0, + regex_max_attempts: int = 10, + sampling_params: Optional[Dict[str, Any]] = None, + init_kwargs: Optional[Dict[str, Any]] = None, + output_directory: Optional[str] = None, + ) -> None: + self.llm_id = llm_id + self.model = model + key = api_key or os.getenv("OPENAI_API_KEY") + if not key: + raise RuntimeError( + "Set OPENAI_API_KEY as global environment variable or pass api_key." + ) + client_kwargs: Dict[str, Any] = {"api_key": key, "timeout": timeout_s} + if base_url: + client_kwargs["base_url"] = base_url + self.client = AsyncOpenAI(**client_kwargs) + + # Sampling/default request params set at init + self.sampling_params = sampling_params + self.use_reasoning = model in reasoning_models + if self.use_reasoning: + self.sampling_params["reasoning"] = { + "effort": "low", + "summary": "detailed", + } + self.regex_max_attempts = max(1, int(regex_max_attempts)) + + def get_inference_policies(self) -> Dict[str, Callable]: + return { + self.llm_id: self.get_action, + } + + async def prepare_adapter_for_inference(self, *args: Any, **kwargs: Any) -> None: + await asyncio.sleep(0) + pass + + async def toggle_eval_mode(self, *args: Any, **kwargs: Any) -> None: + await asyncio.sleep(0) + pass + + async def toggle_training_mode(self, *args: Any, **kwargs: Any) -> None: + await asyncio.sleep(0) + pass + + async def export_adapters(self, *args: Any, **kwargs: Any) -> None: + await asyncio.sleep(0) + pass + + async def checkpoint_all_adapters(self, *args: Any, **kwargs: Any) -> None: + await asyncio.sleep(0) + pass + + def extract_output_from_response(self, resp: Response) -> LLMInferenceOutput: + if len(resp.output) > 1: + summary = resp.output[0].summary + if summary != []: + reasoning_content = summary[0].text + reasoning_content = f"OpenAI Reasoning Summary: {reasoning_content}" + else: + reasoning_content = None + content = resp.output[1].content[0].text + else: + reasoning_content = None + content = resp.output[0].content[0].text + + return LLMInferenceOutput( + content=content, + reasoning_content=reasoning_content, + ) + + @backoff.on_exception( + backoff.expo, Exception, max_time=10**10, max_tries=10**10 + ) + async def get_action( + self, + state: list[ChatTurn], + agent_id: str, + regex: Optional[str] = None, + ) -> LLMInferenceOutput: + # Remove any non-role/content keys from the prompt else openai will error + + # TODO: + prompt = [{"role": p.role, "content": p.content} for p in state] + + # if self.sleep_between_requests: + # await self.wait_random_time() + + # If regex is required, prime the model and validate client-side + if regex: + constraint_msg = { + "role": "user", + "content": ( + f"Output must match this regex exactly: {regex} \n" + "Return only the matching string, with no quotes or extra text." + ), + } + prompt = [constraint_msg, *prompt] + pattern = re.compile(regex) + for _ in range(self.regex_max_attempts): + resp = await self.client.responses.create( + model=self.model, + input=prompt, + **self.sampling_params, + ) + policy_output = self.extract_output_from_response(resp) + if pattern.fullmatch(policy_output.content): + return policy_output + prompt = [ + *prompt, + { + "role": "user", + "content": ( + f"Invalid response format. Expected format (regex): {regex}\n Please try again and provide ONLY a response that matches this regex." + ), + }, + ] + return policy_output + + # Simple, unconstrained generation + resp = await self.client.responses.create( + model=self.model, + input=prompt, + **self.sampling_params, + ) + policy_output = self.extract_output_from_response(resp) + return policy_output + + def shutdown(self) -> None: + self.client = None diff --git a/src_code_for_reproducibility/models/large_language_model_local.py b/src_code_for_reproducibility/models/large_language_model_local.py new file mode 100644 index 0000000000000000000000000000000000000000..7eac1c32c0233cf04106fa12be333ebf74319c2a --- /dev/null +++ b/src_code_for_reproducibility/models/large_language_model_local.py @@ -0,0 +1,384 @@ +""" +TODO: Figure out how to tweak SGlang not to go OOM when batch size is 32. See https://github.com/sgl-project/sglang/issues/6309. +""" + +import logging +import os +import re +import sys +import uuid +from collections.abc import Callable +from copy import deepcopy +from datetime import datetime +from typing import Literal + +import httpx +import requests +import torch +import torch.nn as nn + +# from sglang.utils import ( +# launch_server_cmd, +# print_highlight, +# terminate_process, +# wait_for_server, +# ) +from torch.optim import SGD, Adam, AdamW, RMSprop +from transformers import AutoModelForCausalLM, AutoTokenizer +from trl import AutoModelForCausalLMWithValueHead + +from mllm.chat_utils.apply_template import chat_turns_to_token_ids +from mllm.markov_games.rollout_tree import ChatTurn +from mllm.models.adapter_training_wrapper import AdapterWrapper +from mllm.models.inference_backend import LLMInferenceOutput +from mllm.models.inference_backend_dummy import DummyInferenceBackend +from mllm.models.inference_backend_sglang import SGLangOfflineBackend +from mllm.models.inference_backend_vllm import VLLMAsyncBackend + +logger = logging.getLogger(__name__) +logger.addHandler(logging.StreamHandler(sys.stdout)) + +AdapterID = str +PolicyID = str + + +class LeanLocalLLM: + """ + TOWRITE + """ + + def __init__( + self, + llm_id: str = "base_llm", + model_name: str = "Qwen/Qwen3-4B-Instruct-2507", + device: str = "cuda", + hf_kwargs: dict = {}, + adapter_configs: dict = {}, + output_directory: str = "./models/", + inference_backend: Literal["vllm", "sglang", "dummy"] = "vllm", + inference_backend_sampling_params: dict = {}, + inference_backend_init_kwargs: dict = {}, + initial_adapter_paths: dict[str, str] | None = None, + initial_buffer_paths: list[str] | None = None, + enable_thinking: bool = None, + regex_max_attempts: int = -1, + max_thinking_characters: int = 0, + ): + self.inference_backend_name = inference_backend + self.output_directory = output_directory + self.llm_id = llm_id + self.device = torch.device(device) if device else torch.device("cuda") + self.model_name = model_name + self.adapter_configs = adapter_configs + self.adapter_ids = list(adapter_configs.keys()) + self.enable_thinking = enable_thinking + self.regex_max_attempts = regex_max_attempts + self.initial_buffer_paths = initial_buffer_paths + self.max_thinking_characters = max_thinking_characters + self.regex_retries_count = 0 + + # Optional user-specified initial adapter weight locations (local or HF Hub) + # Format: {adapter_id: path_or_repo_id} + self.initial_adapter_paths: dict[str, str] | None = initial_adapter_paths + + # Path management / imports + self.save_path = str(os.path.join(output_directory, model_name, "adapters")) + self.adapter_paths = { + adapter_id: os.path.join(self.save_path, adapter_id) + for adapter_id in self.adapter_ids + } + checkpoints_dir = os.path.join(self.output_directory, "checkpoints") + self.past_agent_adapter_paths = {} + if os.path.isdir(checkpoints_dir): + for dirname in os.listdir(checkpoints_dir): + dirpath = os.path.join(checkpoints_dir, dirname) + if os.path.isdir(dirpath): + self.past_agent_adapter_paths[f"{dirname}_buffer"] = os.path.join( + dirpath, "agent_adapter" + ) + logger.info( + f"Loaded {len(self.past_agent_adapter_paths)} past agent adapters from checkpoints directory." + ) + if self.initial_buffer_paths is not None: + previous_count = len(self.past_agent_adapter_paths) + for path in self.initial_buffer_paths: + if os.path.isdir(path): + for dirname in os.listdir(path): + dirpath = os.path.join(path, dirname) + if os.path.isdir(dirpath): + self.past_agent_adapter_paths[ + f"{dirname}_buffer" + ] = os.path.join(dirpath, "agent_adapter") + else: + logger.warning( + f"Initial buffer path {path} does not exist or is not a directory." + ) + logger.info( + f"Loaded {len(self.past_agent_adapter_paths) - previous_count} past agent adapters from user-specified initial buffer paths." + ) + self.past_agent_adapter_ids = list(self.past_agent_adapter_paths.keys()) + + # ID management for tracking adapter versions + self.adapter_train_ids = { + adapter_id: self.short_id_generator() for adapter_id in self.adapter_ids + } + # Initialize tokenizer + self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) + # Setup padding token to be same as EOS token + self.tokenizer.pad_token_id = self.tokenizer.eos_token_id + self.tokenizer.pad_token = self.tokenizer.eos_token + + self.weights_got_updated: dict[AdapterID, bool] = { + adapter_id: False for adapter_id in self.adapter_ids + } + self.weights_got_updated.update( + {adapter_id: False for adapter_id in self.past_agent_adapter_ids} + ) + self.current_lora_request = None + self.currently_loaded_adapter_id = None + + # --------------------------------------------------------- + # Init HF model, peft adapters + # --------------------------------------------------------- + self.shared_hf_llm = AutoModelForCausalLM.from_pretrained( + pretrained_model_name_or_path=model_name, + **hf_kwargs, + ) + self.hf_adapters = {} + self.optimizers = {} + for adapter_id in self.adapter_ids: + # Prefer output-folder path if it exists; else fall back to user-specified initial path if provided + output_path = os.path.join(self.save_path, adapter_id) + chosen_path: str | None = None + if os.path.isdir(output_path) and os.listdir(output_path): + chosen_path = output_path + logger.info( + f"Initializing adapter '{adapter_id}': using existing weights from output folder '{chosen_path}'." + ) + elif ( + self.initial_adapter_paths and adapter_id in self.initial_adapter_paths + ): + chosen_path = self.initial_adapter_paths[adapter_id] + logger.info( + f"Initializing adapter '{adapter_id}': using provided initial path '{chosen_path}'." + ) + else: + logger.info( + f"Initializing adapter '{adapter_id}': no initial weights provided or found; starting from scratch." + ) + hf_adapter = AdapterWrapper( + shared_llm=self.shared_hf_llm, + adapter_id=adapter_id, + lora_config=adapter_configs[adapter_id], + path=chosen_path, + ).to(device) + self.hf_adapters[adapter_id] = hf_adapter + # Persist current state of all adapters (ensures remote loads are cached to disk) + self.export_adapters() + + # --------------------------------------------------------- + # Init inference inference_backend + # --------------------------------------------------------- + + if inference_backend == "sglang": + self.inference_backend = SGLangOfflineBackend( + model_name=self.model_name, + save_path=self.save_path, + adapter_paths=self.adapter_paths, + tokenizer=self.tokenizer, + kwargs=inference_backend_init_kwargs, + ) + elif inference_backend == "vllm": + self.inference_backend = VLLMAsyncBackend( + model_name=self.model_name, + # adapter_paths=self.adapter_paths, + tokenizer=self.tokenizer, + engine_init_kwargs=inference_backend_init_kwargs, + sampling_params=inference_backend_sampling_params, + ) + elif inference_backend == "dummy": + self.inference_backend = DummyInferenceBackend() + else: + raise ValueError(f"Unknown inference_backend: {inference_backend}") + + def reset_regex_retries_count(self) -> None: + self.regex_retries_count = 0 + + def get_inference_policies(self) -> dict[PolicyID, Callable]: + """ + TOWRITE + """ + policies = {} + for adapter_id in self.adapter_ids: + # define policy func + async def policy( + state: list[ChatTurn], + agent_id: str, + regex: str | None = None, + _adapter_id=adapter_id, + ): + self.prepare_adapter_for_inference(adapter_id=_adapter_id) + response = await self.get_action(state, agent_id, regex) + return response + + policies[self.llm_id + "/" + adapter_id] = policy + + for adapter_id in self.past_agent_adapter_ids: + # define policy func + async def policy( + state: list[ChatTurn], + agent_id: str, + regex: str | None = None, + _adapter_id=adapter_id, + ): + self.prepare_adapter_for_inference(adapter_id=_adapter_id) + response = await self.get_action(state, agent_id, regex) + return response + + policies[self.llm_id + "/" + adapter_id] = policy + return policies + + def get_adapter_modules(self) -> dict[PolicyID, nn.Module]: + """ + Returns wrappers over the adapters which allows them be + interfaced like regular PyTorch models. + # TODO: create the adapter wrappers here + See adapter_wrapper.py + """ + trainable_objects = {an: self.hf_adapters[an] for an in self.adapter_ids} + return trainable_objects + + async def toggle_training_mode(self) -> None: + for adn in self.adapter_ids: + self.adapter_train_ids[adn] = self.short_id_generator() + await self.inference_backend.toggle_training_mode() + + async def toggle_eval_mode(self) -> None: + await self.inference_backend.toggle_eval_mode() + + def prepare_adapter_for_inference(self, adapter_id: AdapterID) -> None: + self.inference_backend.prepare_adapter( + adapter_id, + adapter_path=self.adapter_paths.get( + adapter_id, self.past_agent_adapter_paths.get(adapter_id, None) + ), + weights_got_updated=self.weights_got_updated[adapter_id], + ) + self.currently_loaded_adapter_id = adapter_id + self.weights_got_updated[adapter_id] = False + + # def _make_prompt_text(self, prompt: list[dict]) -> str: + # if self.enable_thinking is not None: + # prompt_text = self.tokenizer.apply_chat_template( + # prompt, + # tokenize=False, + # add_generation_prompt=True, + # enable_thinking=self.enable_thinking, + # ) + # else: + # prompt_text = self.tokenizer.apply_chat_template( + # prompt, + # tokenize=False, + # add_generation_prompt=True, + # ) + + # return prompt_text + + async def get_action( + self, state: list[ChatTurn], agent_id: str, regex: str | None = None + ) -> ChatTurn: + current_regex = regex if self.regex_max_attempts == -1 else None + pattern = re.compile(regex) if regex else None + nb_attempts = 0 + state = state[:] + while True: + context_token_ids = chat_turns_to_token_ids( + chats=state, + tokenizer=self.tokenizer, + enable_thinking=self.enable_thinking, + ) + # print(f"context is {self.tokenizer.decode(context_token_ids)}") + policy_output = await self.inference_backend.generate( + input_token_ids=context_token_ids.tolist(), + extract_thinking=(self.max_thinking_characters > 0), + regex=current_regex, + ) + # print(f"generated: {self.tokenizer.decode(policy_output.out_token_ids)}") + if ( + pattern is None + or (pattern.fullmatch(policy_output.content)) + or (nb_attempts >= self.regex_max_attempts) + ): + return ChatTurn( + agent_id=agent_id, + role="assistant", + content=policy_output.content, + reasoning_content=policy_output.reasoning_content, + out_token_ids=policy_output.out_token_ids, + log_probs=policy_output.log_probs, + is_state_end=False, + ) + else: + self.regex_retries_count += 1 + nb_attempts += 1 + logger.warning( + f"Response {policy_output.content} did not match regex: {regex}, retry {nb_attempts}/{self.regex_max_attempts}" + ) + if nb_attempts == self.regex_max_attempts: + current_regex = regex + # regex_prompt = ChatTurn( + # role="user", + # content=f"Invalid response format. Expected format (regex): {current_regex}\n Please try again and provide ONLY a response that matches this regex.", + # reasoning_content=None, + # log_probs=None, + # out_token_ids=None, + # is_state_end=False, + # ) + # state.append(regex_prompt) + + def export_adapters(self) -> None: + """ + Any peft wrapper, by default, saves all adapters, not just the one currently loaded. + """ + + # New version of the adapters available + for adapter_id in self.adapter_ids: + self.weights_got_updated[adapter_id] = True + for adapter_id in self.past_agent_adapter_ids: + self.weights_got_updated[adapter_id] = True + + # import random + # self.save_path = self.save_path + str(random.randint(1,500)) + # print(f"Save path: {self.save_path}") + # self.adapter_paths = {adapter_id:os.path.join(self.save_path, adapter_id) for adapter_id in self.adapter_ids} + + adapter_id = self.adapter_ids[0] + self.hf_adapters[adapter_id].save_pretrained(self.save_path) + + def checkpoint_all_adapters(self, checkpoint_indicator: str) -> None: + """ + Checkpoints all adapters to the configured output directory. + """ + adapter_id = self.adapter_ids[0] + output_dir = os.path.join(self.output_directory, "checkpoints") + os.makedirs(output_dir, exist_ok=True) + date_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + agent_adapter_dir = f"{adapter_id}-{checkpoint_indicator}-{date_str}" + export_path = os.path.join(output_dir, agent_adapter_dir) + for adapter_id in self.adapter_ids: + if "agent" in adapter_id: + self.past_agent_adapter_paths[ + f"{agent_adapter_dir}_buffer" + ] = os.path.join(export_path, adapter_id) + self.past_agent_adapter_ids.append(f"{agent_adapter_dir}_buffer") + self.weights_got_updated[f"{agent_adapter_dir}_buffer"] = False + self.hf_adapters[adapter_id].save_pretrained(export_path) + + def short_id_generator(self) -> str: + """ + Generates a short unique ID for tracking adapter versions. + + Returns: + int: An 8-digit integer ID. + """ + return str(uuid.uuid4().int)[:8] diff --git a/src_code_for_reproducibility/models/scalar_critic.py b/src_code_for_reproducibility/models/scalar_critic.py new file mode 100644 index 0000000000000000000000000000000000000000..b0cabf6acfb6757db2871778a397bdbe38b813dd --- /dev/null +++ b/src_code_for_reproducibility/models/scalar_critic.py @@ -0,0 +1,54 @@ +import torch, torch.nn as nn, torch.optim as optim +from transformers import AutoModelForCausalLM, AutoTokenizer +from peft import LoraConfig, get_peft_model + +from mllm.models.adapter_training_wrapper import AdapterWrapper + + +class ScalarCritic(nn.Module): + """ + A causal-LM critic_adapter + a scalar value head: + V_Ο†(s) = wα΅€ h_last + b + Only LoRA adapters (inside critic_adapter) and the value head are trainable. + """ + def __init__(self, critic_adapter: AdapterWrapper): + super().__init__() + self.critic_adapter = critic_adapter + hidden_size = self.critic_adapter.shared_llm.config.hidden_size + self.value_head = nn.Linear(hidden_size, 1).to( + dtype=critic_adapter.dtype, + device=critic_adapter.device) + + def forward(self, + input_ids, + attention_mask=None, + **kwargs): + # AdapterWrapper activates its own adapter internally + outputs = self.critic_adapter( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + **kwargs, + ) + h_last = outputs.hidden_states[-1] # (B, S, H) + values = self.value_head(h_last).squeeze(-1) # (B, S) + return values + + def parameters(self, recurse: bool = True): + """Iterator over *trainable* parameters for this critic.""" + # 1) LoRA params for *this* adapter + for p in self.critic_adapter.parameters(): + yield p + # 2) scalar head + yield from self.value_head.parameters() + + def gradient_checkpointing_enable(self, *args, **kwargs): + self.critic_adapter.gradient_checkpointing_enable(*args, **kwargs) + + @property + def dtype(self): + return self.critic_adapter.dtype + + @property + def device(self): + return self.critic_adapter.device diff --git a/src_code_for_reproducibility/training/README.md b/src_code_for_reproducibility/training/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e1d56e2e3d2fb427f8da2488de4dd9de89650ca2 --- /dev/null +++ b/src_code_for_reproducibility/training/README.md @@ -0,0 +1,20 @@ +Suppose we have a trajectory with 3 timesteps. +token: "0 1 2 3 4 5 6 7 8 9 . . . . ." +string: "A B C a b c A a A a b c A B C" (Capitalized = User, Lowercased = Assistant) +action_mask: "x x x βœ“ βœ“ βœ“ x βœ“ x βœ“ βœ“ βœ“ x x x" (F = False, T = True) +rewards: "r r r r r r R R R R R R r r r" +timestep: "0 0 0 0 0 0 1 1 1 1 1 1 2 2 2" +state_ends: "x x βœ“ x x x βœ“ x x x x x x x βœ“" + +There must be one baseline flag per timestep! + +Then, we might have + +A naive way to interpret this is to think of the number of assistant messages as the number of +steps in the environment. However, this is not the case in practice. Indeed, in a +single simulation step, + + + + +A subtlety arises with credit assignment. In the multi-agent case, we might diff --git a/src_code_for_reproducibility/training/__init__.py b/src_code_for_reproducibility/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src_code_for_reproducibility/training/__pycache__/__init__.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d68eb043e149c6005bbcdfe2fbfb1e1d1d61819 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/__init__.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/annealing_methods.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/annealing_methods.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73991155f0a6d13f03b43dfcd777a49610c88c0e Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/annealing_methods.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/credit_methods.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/credit_methods.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..549ee7cccb7bf669178129df112ae091e2389518 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/credit_methods.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/produce_training_stats.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/produce_training_stats.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eeb7e4551d233ca00f0346838c87e7abbdb84fa1 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/produce_training_stats.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/tally_basic.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tally_basic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96098409cb3e042778965df4c94f7b2a45a66101 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tally_basic.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/tally_metrics.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tally_metrics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c2ee27df0f569f1b270d0b7e2977afa9c954181 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tally_metrics.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/tally_rollout.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tally_rollout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c27a1b8c25697b30b7642894b4ac32e5f5b41531 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tally_rollout.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/tally_tokenwise.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tally_tokenwise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..afe36bd9eb6fa41bb714296a853748846986ed3c Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tally_tokenwise.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/tokenize_chats.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tokenize_chats.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb3cd1234f6a70b27bef4626f66ec30d5ab52517 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tokenize_chats.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/trainer_ad_align.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/trainer_ad_align.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a1cc78df9edd45b4e52fc1f2685eb3a99d48d23 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/trainer_ad_align.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/trainer_common.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/trainer_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cc7e9a6c2510ae34731aa3d46ca42829a483bba4 Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/trainer_common.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/trainer_independent.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/trainer_independent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45d61bbfffaa2299ef7c915156ead08105ef64af Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/trainer_independent.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/trainer_sum_rewards.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/trainer_sum_rewards.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad600757f4ff3da914a94f33768dba90070148cb Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/trainer_sum_rewards.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/__pycache__/training_data_utils.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/training_data_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..01210498226a9a10c4c6675edd9108c376cad74b Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/training_data_utils.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/training/annealing_methods.py b/src_code_for_reproducibility/training/annealing_methods.py new file mode 100644 index 0000000000000000000000000000000000000000..1d46d6fe04d9482ab4f2c5763f4735390d7ecdf6 --- /dev/null +++ b/src_code_for_reproducibility/training/annealing_methods.py @@ -0,0 +1,6 @@ +import numpy as np + + +def sigmoid_annealing(step: int, temperature: float) -> float: + return 2 / (1 + np.exp(-step / temperature)) - 1 + diff --git a/src_code_for_reproducibility/training/tally_metrics.py b/src_code_for_reproducibility/training/tally_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..9d3675c78a66b65a7d538975586c5882ee11b84c --- /dev/null +++ b/src_code_for_reproducibility/training/tally_metrics.py @@ -0,0 +1,55 @@ +import os +from numbers import Number +from typing import Union + +import wandb + + +class Tally: + """ + Minimal scalar-first tally. + - Keys are strings. + - First add stores a scalar; subsequent adds upgrade to a list of scalars. + """ + + def __init__(self): + self.stats = {} + + def reset(self): + self.stats = {} + + def _coerce_scalar(self, value: Union[int, float]) -> Union[int, float]: + if hasattr(value, "item") and callable(getattr(value, "item")): + try: + value = value.item() + except Exception: + pass + if isinstance(value, Number): + return value + raise AssertionError("Metric must be a scalar number") + + def add_metric(self, path: str, metric: Union[int, float]): + metric = float(metric) + assert isinstance(path, str), "Path must be a string." + assert isinstance(metric, float), "Metric must be a scalar number." + + scalar = self._coerce_scalar(metric) + existing = self.stats.get(path) + if existing is None: + self.stats[path] = scalar + elif isinstance(existing, list): + existing.append(scalar) + else: + self.stats[path] = [existing, scalar] + + def save(self, identifier: str, folder: str): + os.makedirs(name=folder, exist_ok=True) + try: + import pickle + + pkl_path = os.path.join(folder, f"{identifier}.tally.pkl") + payload = self.stats + with open(pkl_path, "wb") as f: + pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL) + except Exception: + pass diff --git a/src_code_for_reproducibility/training/tally_rollout.py b/src_code_for_reproducibility/training/tally_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..18bbff345510b55f36aafc88f3f8be9ef46657cc --- /dev/null +++ b/src_code_for_reproducibility/training/tally_rollout.py @@ -0,0 +1,137 @@ +import json +import os +from copy import deepcopy +from typing import Union + +import numpy as np +import pandas as pd +import torch +from transformers import AutoTokenizer + + +class RolloutTallyItem: + def __init__(self, crn_ids: list[str], rollout_ids: list[str], agent_ids: list[str], metric_matrix: torch.Tensor): + """ + Initializes the RolloutTallyItem object. + + Args: + crn_ids (list[str]): List of CRN IDs. + rollout_ids (list[str]): List of rollout IDs. + agent_ids (list[str]): List of agent IDs. + metric_matrix (torch.Tensor): Metric matrix. + """ + if isinstance(crn_ids, torch.Tensor): + crn_ids = crn_ids.detach().cpu().numpy() + if isinstance(rollout_ids, torch.Tensor): + rollout_ids = rollout_ids.detach().cpu().numpy() + if isinstance(agent_ids, torch.Tensor): + agent_ids = agent_ids.detach().cpu().numpy() + self.crn_ids = crn_ids + self.rollout_ids = rollout_ids + self.agent_ids = agent_ids + metric_matrix = metric_matrix.detach().cpu() + assert 0 < metric_matrix.ndim <= 2, "Metric matrix must have less than or equal to 2 dimensions" + if metric_matrix.ndim == 1: + metric_matrix = metric_matrix.reshape(1, -1) + # Convert to float32 if tensor is in BFloat16 format (not supported by numpy) + if metric_matrix.dtype == torch.bfloat16: + metric_matrix = metric_matrix.float() + self.metric_matrix = metric_matrix.numpy() + +class RolloutTally: + """ + Tally is a utility class for collecting and storing training metrics. + It supports adding metrics at specified paths and saving them to disk. + """ + + def __init__(self): + """ + Initializes the RolloutTally object. + + Args: + tokenizer (AutoTokenizer): Tokenizer for converting token IDs to strings. + max_context_length (int, optional): Maximum context length for contextualized metrics. Defaults to 30. + """ + # Array-preserving structure (leaf lists hold numpy arrays / scalars) + self.metrics = {} + # Global ordered list of sample identifiers (crn_id, rollout_id) added in the order samples are processed + + def reset(self): + """ + Resets the base and contextualized tallies to empty dictionaries. + """ + self.metrics = {} + + def get_from_nested_dict(self, dictio: dict, path: str): + """ + Retrieves the value at a nested path in a dictionary. + + Args: + dictio (dict): The dictionary to search. + path (list): List of keys representing the path. + + Returns: + Any: The value at the specified path, or None if not found. + """ + assert isinstance(path, list), "Path must be list." + for sp in path[:-1]: + dictio = dictio.setdefault(sp, {}) + return dictio.get(path[-1], None) + + def set_at_path(self, dictio: dict, path: str, value): + """ + Sets a value at a nested path in a dictionary, creating intermediate dictionaries as needed. + + Args: + dictio (dict): The dictionary to modify. + path (list): List of keys representing the path. + value (Any): The value to set at the specified path. + """ + for sp in path[:-1]: + dictio = dictio.setdefault(sp, {}) + dictio[path[-1]] = value + + + def add_metric( + self, path: list[str], rollout_tally_item: RolloutTallyItem + ): + """ + Adds a metric to the base tally at the specified path. + + Args: + path (list): List of keys representing the path in the base tally. + rollout_tally_item (RolloutTallyItem): The rollout tally item to add. + """ + rollout_tally_item = deepcopy(rollout_tally_item) + + # Update array-preserving tally + array_list = self.get_from_nested_dict(dictio=self.metrics, path=path) + if array_list is None: + self.set_at_path(dictio=self.metrics, path=path, value=[rollout_tally_item]) + else: + array_list.append(rollout_tally_item) + + + def save(self, identifier: str, folder: str): + """ + Saves the base and contextualized tallies to disk as JSON files, and also saves contextualized tallies as CSV files for each game/rollout. + + Args: + path (str): Directory path where the metrics will be saved. + """ + os.makedirs(name=folder, exist_ok=True) + + from datetime import datetime + + now = datetime.now() + + # Pickle only (fastest, exact structure with numpy/scalars at leaves) + try: + import pickle + + pkl_path = os.path.join(folder, f"{identifier}.rt_tally.pkl") + payload = {"metrics": self.metrics} + with open(pkl_path, "wb") as f: + pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL) + except Exception: + pass diff --git a/src_code_for_reproducibility/utils/__init__.py b/src_code_for_reproducibility/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src_code_for_reproducibility/utils/__pycache__/gather_training_stats.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/gather_training_stats.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09bf572c7c52dd7aab1bb6626ad794e15a4b2c55 Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/gather_training_stats.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/__pycache__/kill_sglang.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/kill_sglang.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98167c99cd89bba3ea9b6cb75d28572034a683c9 Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/kill_sglang.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/__pycache__/resource_context.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/resource_context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b9763f1e29c04e7debd256efe7bcfc0cdcfefc8 Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/resource_context.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/__pycache__/rollout_tree_chat_htmls.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_chat_htmls.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33285a6e93f9825a14855114ec1c49eeb89abd74 Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_chat_htmls.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/__pycache__/rollout_tree_gather_utils.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_gather_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fdd056e5ba1aceb1536f2eab49f1c63fe7fd647d Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_gather_utils.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/__pycache__/rollout_tree_stats.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_stats.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5eb56226df7b241e081728dc406c4f0162cba1fa Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_stats.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/__pycache__/stat_pack.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/stat_pack.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65e397d163b96f1a0b90c7541734bd729bb82c9e Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/stat_pack.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/__pycache__/update_start_epoch.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/update_start_epoch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ed6ba62682410f9b2fe69fed3538a96d8fda887 Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/update_start_epoch.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/__pycache__/wandb_utils.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/wandb_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09bd9177e48a4140805ea2695cb567ffe987d70a Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/wandb_utils.cpython-312.pyc differ diff --git a/src_code_for_reproducibility/utils/dict_get_path.py b/src_code_for_reproducibility/utils/dict_get_path.py new file mode 100644 index 0000000000000000000000000000000000000000..cad0acbc4ecd3003dc3ece71a586b4796c0796af --- /dev/null +++ b/src_code_for_reproducibility/utils/dict_get_path.py @@ -0,0 +1,12 @@ + +def get_from_nested_dict(a:dict, path) -> any: + # path is string or list of string + try: + if isinstance(path, str): + return a[path] + else: + for p in path: + a = a[p] + return a + except Exception: + return None diff --git a/src_code_for_reproducibility/utils/format_time.py b/src_code_for_reproducibility/utils/format_time.py new file mode 100644 index 0000000000000000000000000000000000000000..1ad936b01738baaaa9aa03bc8436dcfef1bb4825 --- /dev/null +++ b/src_code_for_reproducibility/utils/format_time.py @@ -0,0 +1,7 @@ +def format_time(seconds): + if seconds >= 3600: + return f"{int(seconds // 3600)}h {int((seconds % 3600) // 60)}m {int(seconds % 60)}s" + elif seconds >= 60: + return f"{int(seconds // 60)}m {int(seconds % 60)}s" + else: + return f"{int(seconds)}s" diff --git a/src_code_for_reproducibility/utils/gather_training_stats.py b/src_code_for_reproducibility/utils/gather_training_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..98ae3d9cad04748059b1b471033178a1f7e4f385 --- /dev/null +++ b/src_code_for_reproducibility/utils/gather_training_stats.py @@ -0,0 +1,257 @@ +import copy +import csv +import gc +import json +import logging +import os +import pickle +import random +import re +import subprocess +import sys +import time +from datetime import datetime +from statistics import mean +from typing import Any, Dict + +import hydra +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch +from omegaconf import OmegaConf + +from mllm.training.tally_metrics import Tally +from mllm.utils.stat_pack import StatPack + + +def get_from_nested_dict(dictio: dict, path: list[str]): + for sp in path[:-1]: + dictio = dictio[sp] + return dictio.get(path[-1]) + + +def set_at_path(dictio: dict, path: list[str], value): + for sp in path[:-1]: + if sp not in dictio: + dictio[sp] = {} + dictio = dictio[sp] + dictio[path[-1]] = value + + +def produce_tabular_render(inpath: str, outpath: str = None): + """ + TODO: docstring + """ + with open(inpath, "r") as f: + data = json.load(f) + rollout_paths = data.keys() + for rollout_path in rollout_paths: + if outpath is None: + m_path = rollout_path.replace("/", "|") + m_path = m_path.replace(".json", "") + m_path = ( + os.path.split(inpath)[0] + + "/contextualized_tabular_renders/" + + m_path + + "_tabular_render.render.csv" + ) + # import pdb; pdb.set_trace() + os.makedirs(os.path.split(m_path)[0], exist_ok=True) + metrics = data[rollout_path] + d = {k: [] for k in metrics[0].keys()} + for m in metrics: + for k, v in m.items(): + d[k].append(v) + d = pd.DataFrame(d) + d.to_csv(m_path) + + +def get_metric_paths(data: list[dict]): + d = data[0] + paths = [] + + def traverse_dict(d, current_path=[]): + for key, value in d.items(): + new_path = current_path + [key] + if isinstance(value, dict): + traverse_dict(value, new_path) + else: + paths.append(new_path) + + traverse_dict(d) + return paths + + +def print_metric_paths(data: list[dict]): + paths = get_metric_paths(data) + for p in paths: + print(p) + + +def get_metric_iteration_list(data: list[dict], metric_path: list[str]): + if isinstance(metric_path, str): + metric_path = [metric_path] + sgl = [] + for d in data: + sgl.append(get_from_nested_dict(d, metric_path)) + return sgl + + +def to_1d_numeric(x): + """Return a 1-D float array (or None if not numeric). Accepts scalars, numpy arrays, or nested list/tuple of them.""" + if x is None: + return None + if isinstance(x, (int, float, np.number)): + return np.array([float(x)], dtype=float) + if isinstance(x, np.ndarray): + try: + return x.astype(float).ravel() + except Exception: + return None + if isinstance(x, (list, tuple)): + parts = [] + for e in x: + arr = to_1d_numeric(e) + if arr is not None and arr.size > 0: + parts.append(arr) + if parts: + return np.concatenate(parts) + return None + return None + + +def get_single_metric_vector(data, metric_path, iterations=None): + if isinstance(metric_path, str): + metric_path = [metric_path] + if iterations == None: + iterations = len(data) + vecs = [] + for d in data: + ar = get_from_nested_dict(d, metric_path) + arr = to_1d_numeric(ar) + if arr is not None: + vecs.append(arr) + + return np.concatenate(vecs) if vecs else np.empty(0, dtype=float) + + +def _load_metrics_file(file_path: str): + if not (file_path.endswith(".tally.pkl") or file_path.endswith(".pkl")): + raise ValueError("Only *.tally.pkl files are supported.") + import pickle + + with open(file_path, "rb") as f: + tree = pickle.load(f) + return tree + + +def get_leaf_items(array_tally: dict, prefix: list[str] = None): + if prefix is None: + prefix = [] + for key, value in array_tally.items(): + next_prefix = prefix + [str(key)] + if isinstance(value, dict): + yield from get_leaf_items(value, next_prefix) + else: + yield next_prefix, value + + +def _sanitize_filename_part(part: str) -> str: + s = part.replace("/", "|") + s = s.replace(" ", "_") + return s + + +def render_rt_tally_pkl_to_csvs(pkl_path: str, outdir: str): + """ + This method takes care of tokenwise logging. + """ + with open(pkl_path, "rb") as f: + payload = pickle.load(f) + # Backward compatibility: older tallies stored the dict directly + if isinstance(payload, dict) and "array_tally" in payload: + array_tally = payload.get("array_tally", {}) + else: + array_tally = payload + + os.makedirs(outdir, exist_ok=True) + trainer_id = os.path.basename(pkl_path).replace(".rt_tally.pkl", "") + for path_list, rollout_tally_items in get_leaf_items(array_tally): + # Create file and initiate writer + path_part = ".".join(_sanitize_filename_part(p) for p in path_list) + filename = f"{trainer_id}__{path_part}.render.csv" + out_path = os.path.join(outdir, filename) + + # Write metric rows to CSV + with open(out_path, "w", newline="") as f: + writer = csv.writer(f) + + # Write header row - need to determine metric column count from first rollout_tally_item + first_item = rollout_tally_items[0] + metric_cols = ( + first_item.metric_matrix.shape[1] + if first_item.metric_matrix.ndim > 1 + else 1 + ) + header = ["agent_id", "crn_id", "rollout_id"] + [ + f"t_{i}" for i in range(metric_cols) + ] + writer.writerow(header) + + for rollout_tally_item in rollout_tally_items: + crn_ids = rollout_tally_item.crn_ids + rollout_ids = rollout_tally_item.rollout_ids + agent_ids = rollout_tally_item.agent_ids + metric_matrix = rollout_tally_item.metric_matrix + for i in range(metric_matrix.shape[0]): + row_vals = metric_matrix[i].reshape(-1) + # Convert row_vals to a list to avoid numpy concatenation issues + row_vals = ( + row_vals.tolist() + if hasattr(row_vals, "tolist") + else list(row_vals) + ) + row_prefix = [ + agent_ids[i], + crn_ids[i], + rollout_ids[i], + ] + writer.writerow(row_prefix + row_vals) + + +def tally_to_stat_pack(tally: Dict[str, Any]): + stat_pack = StatPack() + if "array_tally" in tally: + tally = tally["array_tally"] + + # backward compatibility: will remove later, flatten keys in tally + def get_from_nested_dict(dictio: dict, path: list[str]): + for sp in path[:-1]: + dictio = dictio[sp] + return dictio.get(path[-1]) + + def get_metric_paths(tally: dict): + paths = [] + + def traverse_dict(tally, current_path=[]): + for key, value in tally.items(): + new_path = current_path + [key] + if isinstance(value, dict): + traverse_dict(value, new_path) + else: + paths.append(new_path) + + traverse_dict(tally) + return paths + + paths = get_metric_paths(tally) + modified_tally = {} + for p in paths: + val = get_from_nested_dict(tally, p) + modified_tally["_".join(p)] = np.mean(val) + del tally + tally = modified_tally + for key, value in tally.items(): + stat_pack.add_stat(key, value) + return stat_pack diff --git a/src_code_for_reproducibility/utils/get_coagent_id.py b/src_code_for_reproducibility/utils/get_coagent_id.py new file mode 100644 index 0000000000000000000000000000000000000000..16f0dff5c3f1834aa74c9d8257573b829d62603d --- /dev/null +++ b/src_code_for_reproducibility/utils/get_coagent_id.py @@ -0,0 +1,4 @@ + +def get_coagent_id(ids: list[str], agent_id:str) -> str | None: + for id in ids: + if id != agent_id: return id diff --git a/src_code_for_reproducibility/utils/get_stochastic_game_lengths.py b/src_code_for_reproducibility/utils/get_stochastic_game_lengths.py new file mode 100644 index 0000000000000000000000000000000000000000..a43c386aa1764ae2de1d6e177a0238c633c74bba --- /dev/null +++ b/src_code_for_reproducibility/utils/get_stochastic_game_lengths.py @@ -0,0 +1,30 @@ +import numpy as np + +def get_stochastic_game_lengths( + max_length, + nb_games, + continuation_prob, + same_length_batch=False +): + """ + Generates stochastic game lengths based on a geometric distribution. + + Args: + max_length (int): The maximum length a game can have. + nb_games (int): The number of games to generate lengths for. + continuation_prob (float): The probability of the game continuing after each round. + same_length_batch (bool): If True, all games will have the same length. + + Returns: + Array: An array of game lengths. + """ + if continuation_prob == 1: + return [max_length] * nb_games + if same_length_batch: + length = np.random.geometric(1 - continuation_prob, 1) + game_lengths = np.repeat(length, nb_games) + else: + game_lengths = np.random.geometric(1 - continuation_prob, nb_games) + + game_lengths = np.where(game_lengths > max_length, max_length, game_lengths) + return game_lengths.tolist() diff --git a/src_code_for_reproducibility/utils/kill_sglang.py b/src_code_for_reproducibility/utils/kill_sglang.py new file mode 100644 index 0000000000000000000000000000000000000000..e1e1fe9059e4262e995e0876f7eaec1c7aae4464 --- /dev/null +++ b/src_code_for_reproducibility/utils/kill_sglang.py @@ -0,0 +1,17 @@ +import psutil +import signal + +target_name = "sglang::scheduler" +killed = [] + +def kill_sglang(): + for proc in psutil.process_iter(['pid', 'name', 'cmdline']): + try: + # Some processes may not have a name or cmdline + cmdline = " ".join(proc.info['cmdline']) if proc.info['cmdline'] else "" + if target_name in cmdline: + print(f"Killing PID {proc.pid}: {cmdline}") + proc.send_signal(signal.SIGKILL) + killed.append(proc.pid) + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass diff --git a/src_code_for_reproducibility/utils/output_source_code.py b/src_code_for_reproducibility/utils/output_source_code.py new file mode 100644 index 0000000000000000000000000000000000000000..42b51ecb2b818e5e225af10c58898b8be21dee4d --- /dev/null +++ b/src_code_for_reproducibility/utils/output_source_code.py @@ -0,0 +1,6 @@ +def output_source_code(model, output_path: str) -> None: + """ + Outputs the source code of the model to the given path. + """ + with open(output_path, "w") as f: + f.write(model.source_code) diff --git a/src_code_for_reproducibility/utils/resource_context.py b/src_code_for_reproducibility/utils/resource_context.py new file mode 100644 index 0000000000000000000000000000000000000000..43a3a55d0ca0d4a69eadd0c57650a5afd2ae4831 --- /dev/null +++ b/src_code_for_reproducibility/utils/resource_context.py @@ -0,0 +1,78 @@ +import logging +import time +from contextlib import contextmanager + +import torch + + +def vram_usage(): + output = "" + for i in range(torch.cuda.device_count()): + gpu_memory_allocated = torch.cuda.memory_allocated(i) / ( + 1024**3 + ) # Convert bytes to GB + gpu_memory_reserved = torch.cuda.memory_reserved(i) / ( + 1024**3 + ) # Convert bytes to GB + output += f"GPU {i}: Memory Allocated: {gpu_memory_allocated:.2f} GB, Memory Reserved: {gpu_memory_reserved:.2f} GB" + return output + + +def ram_usage(): + import psutil + + process = psutil.Process() + memory_info = process.memory_info() + ram_used = memory_info.rss / (1024**3) # Convert bytes to GB + return f"RAM Usage: {ram_used:.2f} GB" + + +@contextmanager +def resource_logger_context(logger: logging.Logger, task_description: str): + """ + Context manager to log the resource usage of the current task. + Args: + logger: The logger to use to log the resource usage. + task_description: The description of the task to log. + Returns: + None + """ + try: + initial_time = time.time() + # Assume CUDA is available and use device 0 only + total_mem_bytes = torch.cuda.get_device_properties(0).total_memory + initial_total_bytes = ( + torch.cuda.memory_allocated(0) + torch.cuda.memory_reserved(0) + ) + torch.cuda.reset_peak_memory_stats(0) + yield None + finally: + final_time = time.time() + # Ensure kernels within the block are accounted for + torch.cuda.synchronize() + + # Compute metrics + final_allocated_bytes = torch.cuda.memory_allocated(0) + final_reserved_bytes = torch.cuda.memory_reserved(0) + final_total_bytes = final_allocated_bytes + final_reserved_bytes + + delta_vram_percent_total = ( + 100 * (final_total_bytes - initial_total_bytes) / total_mem_bytes + if total_mem_bytes + else 0.0 + ) + current_percent_vram_taken = ( + 100 * final_total_bytes / total_mem_bytes if total_mem_bytes else 0.0 + ) + block_peak_percent = ( + 100 * torch.cuda.max_memory_allocated(0) / total_mem_bytes + if total_mem_bytes + else 0.0 + ) + delta_time_str = time.strftime( + '%H:%M:%S', time.gmtime(final_time - initial_time) + ) + + logger.info( + f"For task: {task_description}, Ξ”VRAM % (total): {delta_vram_percent_total:.2f}%, Current % of VRAM taken: {current_percent_vram_taken:.2f}%, Block Peak % of device VRAM: {block_peak_percent:.2f}%, Ξ”Time: {delta_time_str}" + ) diff --git a/src_code_for_reproducibility/utils/rollout_tree_chat_htmls.py b/src_code_for_reproducibility/utils/rollout_tree_chat_htmls.py new file mode 100644 index 0000000000000000000000000000000000000000..38971e586269756e2cee8516506b67cf462c2a2d --- /dev/null +++ b/src_code_for_reproducibility/utils/rollout_tree_chat_htmls.py @@ -0,0 +1,664 @@ +from pathlib import Path +from typing import List + +from mllm.utils.rollout_tree_gather_utils import * + + +def html_from_chat_turns(chat_turns: List[ChatTurnLog]) -> str: + """ + Render chat turns as a single, wrapping sequence of messages in time order. + Keep badge and message bubble styles, include time on every badge and + include rewards on assistant badges. Each message is individually + hide/show by click; when hidden, only the badge remains and "(...)" is + shown inline (not inside a bubble). + """ + import html + import re as _re + + # Prepare ordering: sort by (time_step, original_index) to keep stable order within same step + indexed_turns = list(enumerate(chat_turns)) + indexed_turns.sort(key=lambda t: (t[1].time_step, t[0])) + assistant_agents = sorted({t.agent_id for t in chat_turns if t.role == "assistant"}) + enable_split_view = len(assistant_agents) == 2 + + # CSS styles (simplified layout; no time-step or agent-column backgrounds) + css = """ + + """ + + # HTML structure + html_parts = [ + "", + "", + "", + "", + "Chat Turns", + css, + "", + "", + "", + '
', + '
', + '
', + '', + '', + "timesteps", + '', + '', + '', + '', + "to", + '', + '', + '', + ( + '' + if enable_split_view + else "" + ), + "
", + "
", + '
', + ] + + last_time_step = None + for original_index, turn in indexed_turns: + # Build classes + agent_class = f"agent-{re.sub('[^a-z0-9_-]', '-', turn.agent_id.lower())}" + role_class = f"role-{turn.role}" + collapsed_class = " collapsed" if turn.role == "user" else "" + + # Badge content + if turn.role == "assistant": + name = html.escape(turn.agent_id) + emoji = ' πŸ€–' + raw_val = turn.reward + if isinstance(raw_val, (int, float)): + reward_val = f"{raw_val:.4f}".rstrip("0").rstrip(".") + if len(reward_val) > 8: + reward_val = reward_val[:8] + "…" + else: + reward_val = str(raw_val) + # Format: "πŸ€– Alice β€’ Reward: 5.5556 β€’ πŸ’¬ :" + badge_inner = ( + f'{emoji} {name}' + f' β€’ Reward βš‘ = {reward_val}' + ) + else: + # For user messages, show "User of {Agent ID}" in the badge + name = "User of " + html.escape(turn.agent_id) + emoji = 'βš™οΈ' + # Format (no reward): "βš™οΈ User of Alice β€’ " + badge_inner = f'{emoji} {name} β€’ :' + + badge = f'{badge_inner}' + + # Inline timestep distinction badge at step boundaries (render before first message) + ts_badge_html = "" + if last_time_step is None or turn.time_step != last_time_step: + ts_badge_html = f'⏱ {turn.time_step}' + last_time_step = turn.time_step + + escaped_content = html.escape(turn.content) + reasoning_html = "" + if turn.reasoning_content: + # Normalize reasoning to avoid leading/newline whitespace that creates visual gaps + _raw_reasoning = turn.reasoning_content.replace("\r\n", "\n") + _raw_reasoning = _re.sub( + r"^\s*\n+", "", _raw_reasoning + ) # drop leading blank lines + _raw_reasoning = _re.sub( + r"\*\*(\s*\n\s*)", r"** ", _raw_reasoning + ) # newline right after ** + _raw_reasoning = _re.sub( + r"(\s*\n\s*)\*\*", r" **", _raw_reasoning + ) # newline right before ** + escaped_reasoning = html.escape(_raw_reasoning) + reasoning_html = f'πŸ’­{escaped_reasoning}' + collapsed_text = re.sub(r"\s+", " ", escaped_content).strip() + + html_parts.append( + f'
' + f'
{ts_badge_html}{badge}' + f'{reasoning_html}πŸ’¬ {collapsed_text}' + f'(...)' + f"
" + f"
" + ) + + html_parts.append("
") # close linear flow + if enable_split_view: + import html as _html_mod + + html_parts.append( + '") # flow-split + html_parts.extend(["", ""]) + + return "\n".join(html_parts) + + +def export_html_from_rollout_tree(path: Path, outdir: Path, main_only: bool = False): + """Process a rollout tree file and generate HTML files for each path. + Creates separate HTML files for the main path and each branch path. + The main path is saved in the root output directory, while branch paths + are saved in a 'branches' subdirectory. + + Args: + path: Path to the rollout tree JSON file + outdir: Output directory for HTML files + main_only: If True, only export the main trajectory (default: False) + """ + root = load_rollout_tree(path) + mgid = root.id + + main_path, branch_paths = get_rollout_tree_paths(root) + + outdir.mkdir(parents=True, exist_ok=True) + + # Create branches subdirectory if we have branch paths + if not main_only and branch_paths: + branches_dir = outdir / f"mgid:{mgid}_branches_html_renders" + branches_dir.mkdir(parents=True, exist_ok=True) + + # Generate HTML for the main path + chat_turns = gather_all_chat_turns_for_path(main_path) + html_content = html_from_chat_turns(chat_turns) + output_file = outdir / f"mgid:{mgid}_main_html_render.render.html" + with open(output_file, "w", encoding="utf-8") as f: + f.write(html_content) + + # Generate HTML for each branch path + for path_obj in branch_paths: + chat_turns = gather_all_chat_turns_for_path(path_obj) + + html_content = html_from_chat_turns(chat_turns) + + path_id: str = path_obj.id + output_filename = f"{path_id}_html_render.render.html" + + output_file = branches_dir / output_filename + + with open(output_file, "w", encoding="utf-8") as f: + f.write(html_content) diff --git a/src_code_for_reproducibility/utils/rollout_tree_gather_utils.py b/src_code_for_reproducibility/utils/rollout_tree_gather_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d1844aa5a20c01b591865cad108f1fb1577d3ef4 --- /dev/null +++ b/src_code_for_reproducibility/utils/rollout_tree_gather_utils.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import csv +import os +import pickle +import re +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple + +from mllm.markov_games.rollout_tree import * + + + + + +def load_rollout_tree(path: Path) -> RolloutTreeRootNode: + """Load a rollout tree from a PKL file containing a dict.""" + with open(path, "rb") as f: + data = pickle.load(f) + return RolloutTreeRootNode.model_validate(data) + + +@dataclass +class RolloutNodeList: + id: str + nodes: List[RolloutTreeNode] + + +def get_rollout_tree_paths( + root: RolloutTreeRootNode, mgid: Optional[str] = None +) -> Tuple[RolloutNodeList, List[RolloutNodeList]]: + """ + Returns: + main_path: The main path from the root to the end of the tree. + branch_paths: A list of all branch paths from the root to the end of the tree. + Each branch path contains a list of nodes that are part of the branch, including the nodes from the main path before the branch was taken. + """ + branch_paths = [] + + def collect_path_nodes(current) -> List[RolloutTreeNode]: + """Recursively collect all nodes in a path starting from current node.""" + if current is None: + return [] + + if isinstance(current, RolloutTreeNode): + return [current] + collect_path_nodes(current.child) + + elif isinstance(current, RolloutTreeBranchNode): + # For branch nodes, we only follow the main_child for path collection + if current.main_child: + return [current.main_child] + collect_path_nodes( + current.main_child.child + ) + else: + return [] + + def traverse_for_branches( + current, + main_path_prefix: List[RolloutTreeNode], + path_id: str, + current_time_step: Optional[int] = 0, + ): + """Traverse tree to collect all branch paths.""" + if current is None: + return + + if isinstance(current, RolloutTreeNode): + # Continue traversing with this node added to the main path prefix + new_prefix = main_path_prefix + [current] + traverse_for_branches(current.child, new_prefix, path_id, current.time_step) + + elif isinstance(current, RolloutTreeBranchNode): + # Collect all branch paths + if current.branches: + for agent_id, branch_node_list in current.branches.items(): + if branch_node_list: + # Start with the main path prefix, then recursively collect all nodes in this branch + branch_path_nodes = main_path_prefix.copy() + for branch_node in branch_node_list: + branch_path_nodes.extend(collect_path_nodes(branch_node)) + + # Create proper branch path ID with mgid, agent_id, and time_step + mgid_str = mgid or str(root.id) + branch_path_id = f"mgid:{mgid_str}_type:branch_agent:{agent_id}_time_step:{current_time_step}" + branch_paths.append( + RolloutNodeList(id=branch_path_id, nodes=branch_path_nodes) + ) + + # Process the main child and add to prefix + new_prefix = main_path_prefix + if current.main_child: + new_prefix = main_path_prefix + [current.main_child] + + # Continue traversing the main path + if current.main_child: + traverse_for_branches( + current.main_child.child, + new_prefix, + path_id, + current.main_child.time_step, + ) + + # Collect the main path nodes + main_path_nodes = collect_path_nodes(root.child) + + # Traverse to collect all branch paths + traverse_for_branches(root.child, [], "") + + # Create the main path with proper mgid format + mgid_str = mgid or str(root.id) + main_path = RolloutNodeList(id=f"mgid:{mgid_str}_type:main", nodes=main_path_nodes) + + return main_path, branch_paths + + +class ChatTurnLog(BaseModel): + time_step: int + agent_id: str + role: str + content: str + reasoning_content: Optional[str] = None + is_state_end: bool + reward: float + + +def gather_agent_chat_turns_for_path( + agent_id: str, path: RolloutNodeList +) -> List[ChatTurnLog]: + """Iterate through all chat turns for a specific agent in a path sorted by time step.""" + turns = [] + for node in path.nodes: + action_log = node.step_log.action_logs.get(agent_id, []) + if action_log: + for chat_turn in action_log.chat_turns or []: + turns.append( + ChatTurnLog( + time_step=node.time_step, + agent_id=agent_id, + role=chat_turn.role, + content=chat_turn.content, + reasoning_content=getattr(chat_turn, "reasoning_content", None), + is_state_end=chat_turn.is_state_end, + reward=node.step_log.simulation_step_log.rewards.get( + agent_id, 0 + ), + ) + ) + return turns + + +def gather_all_chat_turns_for_path(path: RolloutNodeList) -> List[ChatTurnLog]: + """Iterate through all chat turns for all agents in a path sorted by time step.""" + turns = [] + + # Collect turns from all agents, but interleave them per timestep by (user, assistant) pairs + for node in path.nodes: + # Build (user[, assistant]) pairs for each agent at this timestep + agent_ids = sorted(list(node.step_log.action_logs.keys())) + per_agent_pairs: Dict[str, List[List[ChatTurnLog]]] = {} + + for agent_id in agent_ids: + action_log = node.step_log.action_logs.get(agent_id) + pairs: List[List[ChatTurnLog]] = [] + current_pair: List[ChatTurnLog] = [] + + if action_log and action_log.chat_turns: + for chat_turn in action_log.chat_turns: + turn_log = ChatTurnLog( + time_step=node.time_step, + agent_id=agent_id, + role=chat_turn.role, + content=chat_turn.content, + reasoning_content=getattr(chat_turn, "reasoning_content", None), + is_state_end=chat_turn.is_state_end, + reward=node.step_log.simulation_step_log.rewards.get( + agent_id, 0 + ), + ) + + if chat_turn.role == "user": + # If a previous pair is open, close it and start a new one + if current_pair: + pairs.append(current_pair) + current_pair = [] + current_pair = [turn_log] + else: + # assistant: attach to an open user message if present; otherwise stand alone + if ( + current_pair + and len(current_pair) == 1 + and current_pair[0].role == "user" + ): + current_pair.append(turn_log) + pairs.append(current_pair) + current_pair = [] + else: + # No preceding user or already paired; treat as its own unit + pairs.append([turn_log]) + + if current_pair: + # Unpaired trailing user message + pairs.append(current_pair) + + per_agent_pairs[agent_id] = pairs + + # Interleave pairs across agents: A1, B1, A2, B2, ... + index = 0 + while True: + added_any = False + for agent_id in agent_ids: + agent_pairs = per_agent_pairs.get(agent_id, []) + if index < len(agent_pairs): + for tl in agent_pairs[index]: + turns.append(tl) + added_any = True + if not added_any: + break + index += 1 + + return turns + + +def chat_turns_to_dict(chat_turns: Iterator[ChatTurnLog]) -> Iterator[Dict[str, Any]]: + """Render all chat turns for a path as structured data for JSON.""" + for chat_turn in chat_turns: + yield chat_turn.model_dump() + + +def get_all_agents(root: RolloutTreeRootNode) -> List[str]: + """list of all agent IDs that appear in the tree.""" + if root.child is None: + return [] + + # Get the first node to extract all agent IDs + first_node = root.child + if isinstance(first_node, RolloutTreeBranchNode): + first_node = first_node.main_child + + if first_node is None: + return [] + + # All agents should be present in the first node + agents = set(first_node.step_log.action_logs.keys()) + agents.update(first_node.step_log.simulation_step_log.rewards.keys()) + + return sorted(list(agents)) + + +def gather_agent_main_rewards(agent_id: str, path: RolloutNodeList) -> List[float]: + """Gather main rewards for a specific agent in a path.""" + rewards = [] + for node in path.nodes: + reward = node.step_log.simulation_step_log.rewards[agent_id] + rewards.append(reward) + return rewards + + +def gather_all_rewards(path: RolloutNodeList) -> List[Dict[AgentId, float]]: + """Gather main rewards from main trajectory in a path.""" + rewards = [] + for node in path.nodes: + rewards.append(node.step_log.simulation_step_log.rewards.copy()) + return rewards + + +def gather_simulation_stats( + path: RolloutNodeList, + filter: Callable[[SimulationStepLog], bool], + stat_func: Callable[[SimulationStepLog], Any], +) -> List[Any]: + """Gather stats from main trajectory in a path.""" + stats = [] + for node in path.nodes: + sl = node.step_log.simulation_step_log + if filter(sl): + stats.append(stat_func(sl)) + return stats + + +def gather_simulation_step_logs(path: RolloutNodeList) -> List[SimulationStepLog]: + """Gather simulation information from main trajectory in a path.""" + infos = [] + for node in path.nodes: + infos.append(node.step_log.simulation_step_log) + return infos + + +def export_chat_logs(path: Path, outdir: Path): + """Process a rollout tree PKL file and generate a JSONL of chat turns as dicts. + Each line contains an object with path_id and chat_turns for a single path. + """ + import json + + root = load_rollout_tree(path) + mgid = root.id + + main_path, branch_paths = get_rollout_tree_paths(root) + all_paths = [main_path] + branch_paths + + outdir.mkdir(parents=True, exist_ok=True) + output_file = outdir / f"mgid:{mgid}_plucked_chats.render.jsonl" + + with open(output_file, "w", encoding="utf-8") as f: + for path_obj in all_paths: + chat_turns = gather_all_chat_turns_for_path(path_obj) + output_obj = { + "path_id": str(path_obj.id), + "chat_turns": list(chat_turns_to_dict(iter(chat_turns))), + } + f.write(json.dumps(output_obj, ensure_ascii=False) + "\n") + + diff --git a/src_code_for_reproducibility/utils/rollout_tree_stats.py b/src_code_for_reproducibility/utils/rollout_tree_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..1ac3cd0e34212e7fdbeba19e501b7d96a5f128e1 --- /dev/null +++ b/src_code_for_reproducibility/utils/rollout_tree_stats.py @@ -0,0 +1,50 @@ +from typing import Any, Callable, List, Tuple + +from mllm.markov_games.rollout_tree import RolloutTreeRootNode +from mllm.markov_games.simulation import SimulationStepLog +from mllm.utils.rollout_tree_gather_utils import ( + gather_simulation_step_logs, + get_rollout_tree_paths, +) +from mllm.utils.stat_pack import StatPack + + +def get_rollout_tree_stat_tally( + rollout_tree: RolloutTreeRootNode, + metrics: List[Callable[[SimulationStepLog], List[Tuple[str, float]]]], +) -> StatPack: + stat_tally = StatPack() + # get simulation step logs + node_list = get_rollout_tree_paths(rollout_tree)[0] + simulation_step_logs = gather_simulation_step_logs(node_list) + for simulation_step_log in simulation_step_logs: + for metric in metrics: + metric_result = metric(simulation_step_log) + if metric_result is not None: + for key, value in metric_result: + stat_tally.add_stat(key, value) + return stat_tally + + +def get_rollout_tree_mean_stats( + rollout_tree: RolloutTreeRootNode, metrics: List[Callable[[SimulationStepLog], Any]] +) -> StatPack: + """Get the mean stats for a rollout tree.""" + stat_tally = get_rollout_tree_stat_tally(rollout_tree, metrics) + return stat_tally.mean() + + +def get_mean_rollout_tree_stats( + rollout_trees: List[RolloutTreeRootNode], + metrics: List[Callable[[SimulationStepLog], Any]], +) -> StatPack: + """Get the mean stats for a list of rollout trees.""" + # TODO complete this + stat_tallies = [ + get_rollout_tree_mean_stats(rollout_tree, metrics) + for rollout_tree in rollout_trees + ] + mean_stat_tally = StatPack() + for stat_tally in stat_tallies: + mean_stat_tally.add_stats(stat_tally) + return mean_stat_tally.mean() diff --git a/src_code_for_reproducibility/utils/short_id_gen.py b/src_code_for_reproducibility/utils/short_id_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..bb275cad8fe5e0dd52bebb14e7424e98dcbaba54 --- /dev/null +++ b/src_code_for_reproducibility/utils/short_id_gen.py @@ -0,0 +1,11 @@ +import uuid + + +def generate_short_id() -> int: + """ + Generates a short unique ID for tracking adapter versions. + + Returns: + int: An 8-digit integer ID. + """ + return int(str(uuid.uuid4().int)[:8]) diff --git a/src_code_for_reproducibility/utils/stat_pack.py b/src_code_for_reproducibility/utils/stat_pack.py new file mode 100644 index 0000000000000000000000000000000000000000..46b397139a1a8a4149030a9cc33d2b3afb7b4a12 --- /dev/null +++ b/src_code_for_reproducibility/utils/stat_pack.py @@ -0,0 +1,113 @@ +import csv +import json +import os +import pickle +from collections import Counter +from copy import deepcopy +from locale import strcoll +from statistics import mean +from typing import Any, Dict, Iterator, List, Optional, Tuple, TypedDict + +import matplotlib.pyplot as plt +import numpy as np +from torch.utils.tensorboard import SummaryWriter + +plt.style.use( + "https://raw.githubusercontent.com/dereckpiche/DedeStyle/refs/heads/main/dedestyle.mplstyle" +) + +import wandb + +from . import wandb_utils + + +class StatPack: + def __init__(self): + self.data = {} + + def add_stat(self, key: str, value: float | int | None): + assert ( + isinstance(value, float) or isinstance(value, int) or value is None + ), f"Value {value} is not a valid type" + if key not in self.data: + self.data[key] = [] + self.data[key].append(value) + + def add_stats(self, other: "StatPack"): + for key in other.keys(): + self.add_stat(key, other[key]) + + def __getitem__(self, key: str): + return self.data[key] + + def __setitem__(self, key: str, value: Any): + self.data[key] = value + + def __contains__(self, key: str): + return key in self.data + + def __len__(self): + return len(self.data) + + def __iter__(self): + return iter(self.data) + + def keys(self): + return self.data.keys() + + def values(self): + return self.data.values() + + def items(self): + return self.data.items() + + def mean(self): + mean_st = StatPack() + for key in self.keys(): + if isinstance(self[key], list): + # TODO: exclude None values + non_none_values = [v for v in self[key] if v is not None] + if non_none_values: + mean_st[key] = np.mean(np.array(non_none_values)) + else: + mean_st[key] = None + return mean_st + + def store_plots(self, folder: str): + os.makedirs(folder, exist_ok=True) + for key in self.keys(): + plt.figure(figsize=(10, 5)) + plt.plot(self[key]) + plt.title(key) + plt.savefig(os.path.join(folder, f"{key}.pdf")) + plt.close() + + def store_numpy(self, folder: str): + os.makedirs(folder, exist_ok=True) + for key in self.keys(): + # Sanitize filename components (avoid slashes, spaces, etc.) + safe_key = str(key).replace(os.sep, "_").replace("/", "_").replace(" ", "_") + values = self[key] + # Convert None to NaN for numpy compatibility + arr = np.array( + [(np.nan if (v is None) else v) for v in values], dtype=float + ) + np.save(os.path.join(folder, f"{safe_key}.npy"), arr) + + def store_json(self, folder: str, filename: str = "stats.json"): + os.makedirs(folder, exist_ok=True) + with open(os.path.join(folder, filename), "w") as f: + json.dump(self.data, f, indent=4) + + def store_csv(self, folder: str): + os.makedirs(folder, exist_ok=True) + for key in self.keys(): + with open(os.path.join(folder, f"stats.csv"), "w") as f: + writer = csv.writer(f) + writer.writerow([key] + self[key]) + + def store_pickle(self, folder: str): + os.makedirs(folder, exist_ok=True) + for key in self.keys(): + with open(os.path.join(folder, f"stats.pkl"), "wb") as f: + pickle.dump(self[key], f) diff --git a/src_code_for_reproducibility/utils/update_start_epoch.py b/src_code_for_reproducibility/utils/update_start_epoch.py new file mode 100644 index 0000000000000000000000000000000000000000..036ddce31b12e7a6547c5099dd37962a88055643 --- /dev/null +++ b/src_code_for_reproducibility/utils/update_start_epoch.py @@ -0,0 +1,9 @@ +import os + +# During run, set hydra.run.dir=./outputs/{folder} +def update_start_epoch(cfg, output_directory): + if cfg["experiment"]["resume_experiment"]: + folders = [f for f in os.listdir(output_directory) if f.startswith("iteration_")] + iterations = [int(f.split("_")[1]) for f in folders] if folders else [0] + cfg["experiment"]["start_epoch"] = max(iterations) + return None diff --git a/src_code_for_reproducibility/utils/wandb_utils.py b/src_code_for_reproducibility/utils/wandb_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5d83ed1a5304f208288f78457582b0acdb58c4 --- /dev/null +++ b/src_code_for_reproducibility/utils/wandb_utils.py @@ -0,0 +1,164 @@ +import os +from typing import Any, Dict, Optional + + +_WANDB_AVAILABLE = False +_WANDB_RUN = None + + +def _try_import_wandb(): + global _WANDB_AVAILABLE + if _WANDB_AVAILABLE: + return True + try: + import wandb # type: ignore + + _WANDB_AVAILABLE = True + return True + except Exception: + _WANDB_AVAILABLE = False + return False + + +def _safe_get(cfg: Dict[str, Any], path: list[str], default: Any = None) -> Any: + cur: Any = cfg + for key in path: + if not isinstance(cur, dict) or key not in cur: + return default + cur = cur[key] + return cur + + +def is_enabled(cfg: Dict[str, Any]) -> bool: + return bool(_safe_get(cfg, ["logging", "wandb", "enabled"], False)) + + +def init(cfg: Dict[str, Any], run_dir: str, run_name: Optional[str] = None) -> None: + """ + Initialize Weights & Biases if enabled in config. No-op if disabled or wandb not installed. + """ + global _WANDB_RUN + if not is_enabled(cfg): + return + if not _try_import_wandb(): + return + + import wandb # type: ignore + + project = _safe_get(cfg, ["logging", "wandb", "project"], "llm-negotiation") + entity = _safe_get(cfg, ["logging", "wandb", "entity"], None) + mode = _safe_get(cfg, ["logging", "wandb", "mode"], "online") + tags = _safe_get(cfg, ["logging", "wandb", "tags"], []) or [] + notes = _safe_get(cfg, ["logging", "wandb", "notes"], None) + group = _safe_get(cfg, ["logging", "wandb", "group"], None) + name = _safe_get(cfg, ["logging", "wandb", "name"], run_name) + + # Ensure files are written into the hydra run directory + os.makedirs(run_dir, exist_ok=True) + os.environ.setdefault("WANDB_DIR", run_dir) + + # Convert cfg to plain types for W&B config; fallback to minimal dictionary + try: + from omegaconf import OmegaConf # type: ignore + + cfg_container = OmegaConf.to_container(cfg, resolve=True) # type: ignore + except Exception: + cfg_container = cfg + + _WANDB_RUN = wandb.init( + project=project, + entity=entity, + mode=mode, + name=name, + group=group, + tags=tags, + notes=notes, + config=cfg_container, + dir=run_dir, + reinit=True, + ) + + +def log(metrics: Dict[str, Any], step: Optional[int] = None) -> None: + """Log a flat dictionary of metrics to W&B if active.""" + if not _WANDB_AVAILABLE or _WANDB_RUN is None: + return + try: + import wandb # type: ignore + + wandb.log(metrics if step is None else dict(metrics, step=step)) + except Exception: + pass + + +def _flatten(prefix: str, data: Dict[str, Any], out: Dict[str, Any]) -> None: + for k, v in data.items(): + key = f"{prefix}.{k}" if prefix else k + if isinstance(v, dict): + _flatten(key, v, out) + else: + out[key] = v + + +def _summarize_value(value: Any) -> Dict[str, Any]: + import numpy as np # local import to avoid hard dependency during disabled mode + + if value is None: + return {"none": 1} + # Scalars + if isinstance(value, (int, float)): + return {"value": float(value)} + # Lists or arrays + try: + arr = np.asarray(value) + if arr.size == 0: + return {"size": 0} + return { + "mean": float(np.nanmean(arr)), + "min": float(np.nanmin(arr)), + "max": float(np.nanmax(arr)), + "last": float(arr.reshape(-1)[-1]), + "size": int(arr.size), + } + except Exception: + # Fallback: string repr + return {"text": str(value)} + + +def log_tally(array_tally: Dict[str, Any], prefix: str = "", step: Optional[int] = None) -> None: + """ + Flatten and summarize Tally.array_tally and log to WandB. + Each leaf list/array is summarized with mean/min/max/last/size. + """ + if not _WANDB_AVAILABLE or _WANDB_RUN is None: + return + summarized: Dict[str, Any] = {} + + def walk(node: Any, path: list[str]): + if isinstance(node, dict): + for k, v in node.items(): + walk(v, path + [k]) + return + # node is a list of values accumulated over time + key = ".".join([p for p in ([prefix] if prefix else []) + path]) + try: + summary = _summarize_value(node) + for sk, sv in summary.items(): + summarized[f"{key}.{sk}"] = sv + except Exception: + summarized[f"{key}.error"] = 1 + + walk(array_tally, []) + if summarized: + log(summarized, step=step) + + +def log_flat_stats(stats: Dict[str, Any], prefix: str = "", step: Optional[int] = None) -> None: + if not _WANDB_AVAILABLE or _WANDB_RUN is None: + return + flat: Dict[str, Any] = {} + _flatten(prefix, stats, flat) + if flat: + log(flat, step=step) + +