Add files using upload-large-folder tool

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)

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 "<assistant>{content}</assistant>""
+    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

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 '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if reasoning_content %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\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 %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}
+"""

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)

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

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

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

src_code_for_reproducibility/markov_games/diplomacy/diplomacy_agent.py

@@ -0,0 +1,259 @@
+from typing import Dict, List, Tuple, Optional, Any
+import copy
+
+class DiplomacyAgent:
+    """Agent handler for Diplomacy game that follows the MARL standard.
+    
+    This class is responsible for parsing LLM output into valid Diplomacy orders,
+    managing the agent state, and providing information for logging.
+    """
+    
+    def __init__(self, policy_id: str, power_name: str, random_valid_move=False):
+        """Initialize the agent handler for a power in the Diplomacy game.
+        
+        Args:
+            power_name: The name of the power this agent controls (e.g., 'FRANCE', 'ENGLAND')
+            policy_id: The identifier for the policy this agent uses
+            random_valid_move: If True, will select random valid moves instead of using LLM (default: False)
+        """
+        self.policy_id = policy_id
+        self.power_name = power_name
+        self.orders = []
+        self.wait = True
+        self.processing_state = "WAITING_FOR_ORDERS"
+        self.parsed_orders = []
+        self.order_status = {}
+        self.message_history = []
+        self.random_valid_move = random_valid_move
+        
+    def step(self, observation_from_env, policy_output=None):
+        """Update the agent state based on the observation and LLM output.
+        
+        Args:
+            observation_from_env: The observation from the environment
+            policy_output: The output from the LLM
+            
+        Returns:
+            policy_id: The policy identifier
+            policy_input: The input to the policy
+            action: The official action to be sent to the environment
+            done: Whether the LLM action is ready to be sent to the environment
+            info: Additional information about the agent
+        """
+        info = {}
+        
+        # If random_valid_move is enabled, select random valid moves
+        if self.random_valid_move:
+            valid_orders = self._select_random_valid_moves(observation_from_env)
+            self.orders = valid_orders
+            self.wait = False
+            action = {
+                "orders": valid_orders,
+                "wait": False
+            }
+            return self.policy_id, {}, action, True, info
+        
+        # If no policy output, this is the initial step - prepare prompt
+        if policy_output is None:
+            # Create initial prompt for the LLM
+            phase = observation_from_env.get('phase', '')
+            units = observation_from_env.get('units', {}).get(self.power_name, [])
+            centers = observation_from_env.get('centers', {}).get(self.power_name, [])
+            orderable_locations = observation_from_env.get('orderable_locations', {})
+            
+            prompt = self._create_prompt(phase, units, centers, orderable_locations)
+            
+            return self.policy_id, {"prompt": prompt}, None, False, info
+            
+        # Process the LLM output to extract orders
+        success, parsed_orders = self._parse_llm_output(policy_output)
+        self.parsed_orders = parsed_orders
+        
+        if not success:
+            # Need more information from LLM
+            clarification_prompt = self._create_clarification_prompt(policy_output, parsed_orders)
+            return self.policy_id, {"prompt": clarification_prompt}, None, False, info
+        
+        # Validate if the orders are valid for the current phase
+        valid_orders = self._validate_orders(parsed_orders, observation_from_env)
+        
+        if valid_orders:
+            # Orders are valid, prepare action for environment
+            self.orders = valid_orders
+            self.wait = False
+            action = {
+                "orders": valid_orders,
+                "wait": False
+            }
+            return self.policy_id, {}, action, True, info
+        else:
+            # Orders are invalid, ask for new ones
+            error_prompt = self._create_error_prompt(parsed_orders, observation_from_env)
+            return self.policy_id, {"prompt": error_prompt}, None, False, info
+    
+    def _create_prompt(self, phase, units, centers, orderable_locations):
+        """Create the initial prompt for the LLM.
+        
+        Args:
+            phase: The current game phase
+            units: List of units controlled by this power
+            centers: List of supply centers controlled by this power
+            orderable_locations: List of locations where orders can be issued
+            
+        Returns:
+            A prompt string for the LLM
+        """
+        prompt = f"You are playing as {self.power_name} in Diplomacy. The current phase is {phase}.\n\n"
+        prompt += f"Your units: {', '.join(units)}\n"
+        prompt += f"Your supply centers: {', '.join(centers)}\n"
+        prompt += f"Locations you can order: {', '.join(orderable_locations)}\n\n"
+        
+        if phase.endswith('M'):  # Movement phase
+            prompt += "Please provide orders for your units in the form:\n"
+            prompt += "- A LON H (hold)\n"
+            prompt += "- F NTH - NWY (move)\n"
+            prompt += "- A WAL S F LON (support)\n"
+            prompt += "- F NWG C A NWY - EDI (convoy)\n"
+        elif phase.endswith('R'):  # Retreat phase
+            prompt += "Please provide retreat orders for your dislodged units:\n"
+            prompt += "- A PAR R MAR (retreat to MAR)\n"
+            prompt += "- A PAR D (disband)\n"
+        elif phase.endswith('A'):  # Adjustment phase
+            if len(units) < len(centers):
+                prompt += "You can build units. Please provide build orders:\n"
+                prompt += "- A PAR B (build army in PAR)\n"
+                prompt += "- F BRE B (build fleet in BRE)\n"
+                prompt += "- WAIVE (waive a build)\n"
+            elif len(units) > len(centers):
+                prompt += "You must remove units. Please provide disbandment orders:\n"
+                prompt += "- A PAR D (disband army in PAR)\n"
+                prompt += "- F BRE D (disband fleet in BRE)\n"
+        
+        prompt += "\nProvide your orders as a list, one per line."
+        return prompt
+    
+    def _parse_llm_output(self, llm_output):
+        """Parse the LLM output to extract orders.
+        
+        Args:
+            llm_output: The raw output from the LLM
+            
+        Returns:
+            success: Whether parsing was successful
+            parsed_orders: List of parsed orders
+        """
+        # Simple parsing for now - extract lines that look like orders
+        lines = llm_output.strip().split('\n')
+        orders = []
+        
+        for line in lines:
+            # Remove list markers, hyphens, etc.
+            line = line.strip('- *•').strip()
+            
+            # Skip empty lines and lines that don't look like orders
+            if not line or line.startswith('I ') or line.startswith('Let\'s'):
+                continue
+                
+            # Check if it looks like a Diplomacy order
+            if (' H' in line or ' -' in line or ' S ' in line or ' C ' in line or 
+                ' R ' in line or ' D' in line or ' B' in line or line == 'WAIVE'):
+                orders.append(line)
+        
+        return len(orders) > 0, orders
+    
+    def _validate_orders(self, orders, observation):
+        """Validate if the orders are valid for the current phase.
+        
+        Args:
+            orders: List of orders to validate
+            observation: Current observation from the environment
+            
+        Returns:
+            List of valid orders or None if invalid
+        """
+        # For simplicity, we'll assume all parsed orders are valid
+        # In a real implementation, we would use the game's validation logic
+        return orders
+    
+    def _create_clarification_prompt(self, previous_output, parsed_orders):
+        """Create a prompt asking for clarification when orders couldn't be parsed.
+        
+        Args:
+            previous_output: The previous LLM output
+            parsed_orders: Any orders that were successfully parsed
+            
+        Returns:
+            A prompt string for the LLM
+        """
+        prompt = f"I couldn't fully understand your orders for {self.power_name}. "
+        
+        if parsed_orders:
+            prompt += f"I understood these orders:\n"
+            for order in parsed_orders:
+                prompt += f"- {order}\n"
+                
+        prompt += "\nPlease provide clear, valid Diplomacy orders in the format:\n"
+        prompt += "- A LON H\n- F NTH - NWY\n- etc.\n"
+        return prompt
+    
+    def _create_error_prompt(self, invalid_orders, observation):
+        """Create a prompt when orders are invalid.
+        
+        Args:
+            invalid_orders: The invalid orders
+            observation: Current observation from the environment
+            
+        Returns:
+            A prompt string for the LLM
+        """
+        prompt = f"The following orders for {self.power_name} are invalid:\n"
+        for order in invalid_orders:
+            prompt += f"- {order}\n"
+            
+        prompt += "\nPlease provide valid orders for your units."
+        return prompt
+    
+    def get_log_info(self):
+        """Get information about the agent required to log a trajectory.
+        
+        Returns:
+            log_info: Information about the agent required to log a trajectory.
+        """
+        return {
+            "power_name": self.power_name,
+            "orders": self.orders,
+            "wait": self.wait,
+            "parsing_state": self.processing_state,
+            "message_history": self.message_history
+        }
+    
+    def render(self):
+        """Render the current state of the agent."""
+        print(f"Power: {self.power_name}")
+        print(f"Orders: {self.orders}")
+        print(f"Wait: {self.wait}")
+    
+    def close(self):
+        """Perform any necessary cleanup."""
+        pass
+
+    def _select_random_valid_moves(self, observation):
+        """Select random valid moves for all units.
+        
+        Args:
+            observation: Current observation from the environment
+            
+        Returns:
+            List of valid orders
+        """
+        import random
+        
+        possible_orders = observation.get('possible_orders', {})
+        valid_orders = []
+        
+        # For each location with possible orders, select one randomly
+        for location, orders in possible_orders.items():
+            if orders:  # If there are any possible orders for this location
+                valid_orders.append(random.choice(orders))
+        
+        return valid_orders

src_code_for_reproducibility/markov_games/diplomacy/diplomacy_env.py

@@ -0,0 +1,230 @@
+from typing import Dict, List, Tuple, Optional, Any
+from diplomacy import Game
+import random
+
+class DiplomacyEnv:
+    """Multi-Agent Reinforcement Learning environment for Diplomacy.
+    
+    This class wraps the Diplomacy game engine to provide an interface
+    compliant with the MARL standard.
+    """
+    
+    def __init__(self, random_seed=None, map_name="standard", game_id=None, rules=None, max_steps=50):
+        """Initialize the Diplomacy environment.
+        
+        Args:
+            map_name: The name of the map to use (default: "standard")
+            game_id: Optional game ID
+            rules: Optional rules to apply to the game
+            max_steps: Maximum number of steps before forcing game end (default: 10)
+        """
+        self.random_seed = random_seed
+        self.map_name = map_name
+        self.game_id = game_id
+        self.rules = rules or []
+        self.game = None
+        self.active_powers = []
+        self.render_mode = None
+        self.max_steps = max_steps
+        self.current_steps = 0
+    
+    def reset(self):
+        """Reset the environment to an initial state and return the initial observation.
+        
+        Returns:
+            observation: A dictionary where keys are agent identifiers and values are observations.
+        """
+        # Initialize a new game
+        self.game = Game(game_id=self.game_id, map_name=self.map_name)
+        
+        # Apply rules
+        for rule in self.rules:
+            self.game.add_rule(rule)
+        
+        # Determine active powers (not eliminated)
+        self.active_powers = [name for name, power in self.game.powers.items() 
+                             if not power.is_eliminated()]
+        
+        # Reset step counter
+        self.current_steps = 0
+        
+        # Create initial observations for all powers
+        observations = {}
+        for power_name in self.active_powers:
+            observations[power_name] = self._create_observation(power_name)
+        
+        return observations
+    
+    def step(self, actions):
+        """Take a step in the environment using the provided actions.
+        
+        Args:
+            actions: A dictionary where keys are agent identifiers and values are actions.
+            
+        Returns:
+            observations: A dictionary where keys are agent identifiers and values are observations.
+            done: Whether the episode has ended.
+            info: Additional information about the environment.
+        """
+        print(f"stepping {self.current_steps}")
+        self.current_steps += 1
+        # Apply actions (orders) for each power
+        for power_name, action in actions.items():
+            if power_name in self.active_powers:
+                orders = action.get("orders", [])
+                wait = action.get("wait", True)
+                
+                # Set orders for the power
+                if orders:
+                    self.game.set_orders(power_name, orders)
+                
+                # Set wait flag
+                self.game.set_wait(power_name, wait)
+        
+        # Check if all active powers are ready to proceed
+        if self.game.does_not_wait():
+            # Process the current phase
+            self.game.process()
+            
+            
+            # Update active powers list after processing
+            self.active_powers = [name for name, power in self.game.powers.items() 
+                                 if not power.is_eliminated()]
+        
+        # Create observations for all active powers
+        observations = {}
+        for power_name in self.active_powers:
+            observations[power_name] = self._create_observation(power_name)
+        
+        # Check if the game is done (either naturally or due to max steps)
+        done = self.game.is_game_done or self.current_steps >= self.max_steps
+        
+        # Create info dict
+        info = {
+            "phase": self.game.get_current_phase(),
+            "active_powers": self.active_powers,
+            "centers": self.game.get_centers(),
+            "units": self.game.get_units(),
+            "current_steps": self.current_steps,
+            "max_steps_reached": self.current_steps >= self.max_steps
+        }
+        
+        return observations, done, info
+    
+    def _create_observation(self, power_name):
+        """Create observation for a specific power.
+        
+        Args:
+            power_name: The name of the power
+            
+        Returns:
+            An observation dictionary
+        """
+        observation = {
+            "phase": self.game.get_current_phase(),
+            "units": self.game.get_units(),
+            "centers": self.game.get_centers(),
+            "orderable_locations": self.game.get_orderable_locations(power_name),
+            "order_status": self.game.get_order_status(power_name),
+            "possible_orders": self._get_possible_orders_for_power(power_name)
+        }
+        return observation
+    
+    def _get_possible_orders_for_power(self, power_name):
+        """Get all possible orders for a power's units.
+        
+        Args:
+            power_name: The name of the power
+            
+        Returns:
+            A dictionary mapping units to their possible orders
+        """
+        all_possible_orders = self.game.get_all_possible_orders()
+        
+        # Filter for only the locations where this power has units
+        power_units = self.game.get_units(power_name)
+        power_unit_locations = [unit[2:] for unit in power_units]
+        
+        # For retreat phases, include retreating units
+        if self.game.phase_type == 'R':
+            power = self.game.get_power(power_name)
+            power_unit_locations.extend([unit[2:] for unit in power.retreats])
+        
+        # For adjustment phases, include buildable locations
+        elif self.game.phase_type == 'A':
+            power = self.game.get_power(power_name)
+            # If we have more centers than units, we can build
+            if len(power.centers) > len(power.units):
+                buildable_sites = self.game._build_sites(power)
+                power_unit_locations.extend(buildable_sites)
+            # If we have more units than centers, we need to remove
+            elif len(power.units) > len(power.centers):
+                # All units are candidates for removal
+                pass
+        
+        # Filter the possible orders to only those for this power's units/locations
+        power_possible_orders = {}
+        for loc, orders in all_possible_orders.items():
+            if loc[:3] in power_unit_locations:
+                power_possible_orders[loc] = orders
+        
+        return power_possible_orders
+    
+    def get_log_info(self):
+        """Get additional information about the environment for logging.
+        
+        Returns:
+            log_info: Information about the environment required to log the game.
+        """
+        if not self.game:
+            return {}
+            
+        return {
+            "game_id": self.game.game_id,
+            "phase": self.game.get_current_phase(),
+            "map_name": self.game.map_name,
+            "centers": self.game.get_centers(),
+            "units": self.game.get_units(),
+            "powers": {name: {
+                "units": power.units,
+                "centers": power.centers,
+                "is_eliminated": power.is_eliminated(),
+                "order_status": self.game.get_order_status(name)
+            } for name, power in self.game.powers.items()},
+            "orders": self.game.get_orders(),
+            "active_powers": self.active_powers,
+            "is_game_done": self.game.is_game_done,
+            "outcome": self.game.outcome if self.game.is_game_done else None
+        }
+    
+    def render(self, mode='human'):
+        """Render the current state of the environment.
+        
+        Args:
+            mode: The rendering mode ('human', 'svg', etc.)
+        
+        Returns:
+            The rendered image if applicable
+        """
+        self.render_mode = mode
+        if self.game:
+            if mode == 'human':
+                # Just print basic game state
+                print(f"Game: {self.game.game_id}")
+                print(f"Phase: {self.game.get_current_phase()}")
+                print(f"Active Powers: {self.active_powers}")
+                print("Supply Centers:")
+                for power_name, centers in self.game.get_centers().items():
+                    print(f"  {power_name}: {centers}")
+                print("Units:")
+                for power_name, units in self.game.get_units().items():
+                    print(f"  {power_name}: {units}")
+                return None
+            elif mode == 'svg':
+                # Return SVG representation
+                return self.game.render(output_format='svg')
+        return None
+    
+    def close(self):
+        """Perform any necessary cleanup."""
+        self.game = None

src_code_for_reproducibility/markov_games/diplomacy/diplomacy_logging.py

@@ -0,0 +1,360 @@
+import os
+import json
+from utils.common_imports import *
+
+
+
+def diplomacy_log_match(
+        path,
+        agents_log_info,
+        env_log_info,
+        metrics_func=None,
+        metrics_func_args=None
+        ):
+    """
+    Logs the Diplomacy game data and generates HTML visualizations using the get_log_info methods.
+    
+    Args:
+        path (str): Base path to save the data.
+        agents_log_info (list): List of agent information dictionaries containing the get_log_info results.
+        env_log_info (dict): Environment information from its get_log_info method.
+        metrics_func (str, optional): Name of the function to calculate metrics.
+        metrics_func_args (dict, optional): Arguments for the metrics function.
+    """
+    # Create directory structure
+    os.makedirs(path, exist_ok=True)
+    
+    # Save the environment log info
+    env_log_path = os.path.join(path, "env_log.json")
+    with open(env_log_path, "w") as f:
+        json.dump(env_log_info, f, indent=4, default=_json_serialize)
+    
+    # Process each agent's log info
+    for agent_log in agents_log_info:
+        power_name = agent_log["power_name"]
+        
+        # Define paths for raw data and statistics subfolders
+        power_path = os.path.join(path, power_name)
+        raw_data_path = os.path.join(power_path, "raw_data")
+        statistics_path = os.path.join(power_path, "statistics")
+        
+        # Ensure directories exist
+        os.makedirs(raw_data_path, exist_ok=True)
+        os.makedirs(statistics_path, exist_ok=True)
+        
+        # Determine the next available file number for raw data
+        raw_files = os.listdir(raw_data_path)
+        raw_numbers = [int(f.split('_')[-1].split('.')[0]) for f in raw_files if f.startswith("log_")]
+        next_raw_number = max(raw_numbers, default=0) + 1
+        raw_file = os.path.join(raw_data_path, f"log_{next_raw_number}.json")
+        
+        # Save agent log info
+        with open(raw_file, "w") as f:
+            json.dump(agent_log, f, indent=4, default=_json_serialize)
+        
+        # Log metrics if a metrics function is provided
+        if metrics_func:
+            metrics_files = os.listdir(statistics_path)
+            metrics_numbers = [int(f.split('_')[-1].split('.')[0]) for f in metrics_files if f.startswith("metrics_")]
+            next_metrics_number = max(metrics_numbers, default=0) + 1
+            metrics_file = os.path.join(statistics_path, f"metrics_{next_metrics_number}.json")
+            
+            metrics = globals()[metrics_func](agent_log, info, **metrics_func_args)
+            with open(metrics_file, "w") as f:
+                json.dump(metrics, f, indent=4)
+    
+    # Generate the HTML visualization
+    html_content = generate_diplomacy_html(agents_log_info, env_log_info)
+    
+    # Ensure the html directory exists
+    html_path = os.path.join(path, "html")
+    os.makedirs(html_path, exist_ok=True)
+    
+    # Determine the next available file number for HTML
+    html_files = os.listdir(html_path)
+    html_numbers = [int(f.split('_')[-1].split('.')[0]) for f in html_files if f.startswith("game_summary_")]
+    next_html_number = max(html_numbers, default=0) + 1
+    html_file = os.path.join(html_path, f"game_summary_{next_html_number}.html")
+    
+    # Save the HTML content to a file
+    with open(html_file, "w") as f:
+        f.write(html_content)
+
+def generate_diplomacy_html(agent_infos, env_info):
+    """
+    Generate HTML visualization for a Diplomacy game.
+    
+    Args:
+        agent_infos (list): List of agent information dictionaries from get_log_info.
+        env_info (dict): Environment information from get_log_info.
+        
+    Returns:
+        str: HTML content for the game visualization.
+    """
+    # Extract game information
+    game_id = env_info.get("game_id", "Unknown")
+    phase = env_info.get("phase", "Unknown")
+    map_name = env_info.get("map_name", "standard")
+    is_game_done = env_info.get("is_game_done", False)
+    outcome = env_info.get("outcome", [])
+    
+    centers = env_info.get("centers", {})
+    units = env_info.get("units", {})
+    
+    # HTML head and style
+    html_content = """
+    <!DOCTYPE html>
+    <html lang="en">
+    <head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
+        <title>Diplomacy Game {game_id}</title>
+        <style>
+            body {{
+                font-family: 'Arial', sans-serif;
+                background-color: #f5f5f5;
+                color: #333333;
+                margin: 0;
+                padding: 20px;
+            }}
+            .container {{
+                display: grid;
+                grid-template-columns: repeat(3, 1fr);
+                grid-gap: 20px;
+                margin-bottom: 30px;
+            }}
+            .central-info {{
+                grid-column: span 3;
+                background: #fff;
+                padding: 20px;
+                border-radius: 10px;
+                box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
+                margin-bottom: 20px;
+            }}
+            .power-column {{
+                background: #fff;
+                padding: 15px;
+                border-radius: 10px;
+                box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
+            }}
+            .message {{
+                margin-bottom: 15px;
+                padding: 12px;
+                border-radius: 8px;
+                box-shadow: 0 1px 4px rgba(0, 0, 0, 0.1);
+            }}
+            .user {{
+                background: rgba(235, 245, 255, 0.8);
+                border-left: 4px solid #007bff;
+            }}
+            .assistant {{
+                background: rgba(240, 255, 240, 0.8);
+                border-right: 4px solid #28a745;
+            }}
+            .orders {{
+                background: rgba(255, 248, 225, 0.8);
+                border-left: 4px solid #ffc107;
+            }}
+            .role {{
+                font-weight: bold;
+                margin-bottom: 5px;
+                color: #333333;
+            }}
+            .power-name {{
+                text-align: center;
+                font-size: 1.4em;
+                margin-bottom: 15px;
+                color: #000;
+                font-weight: 600;
+                text-transform: uppercase;
+                letter-spacing: 1px;
+            }}
+            .game-info {{
+                display: grid;
+                grid-template-columns: repeat(2, 1fr);
+                grid-gap: 15px;
+            }}
+            .info-card {{
+                background: #f9f9f9;
+                padding: 15px;
+                border-radius: 8px;
+                box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
+            }}
+            .supply-centers, .units-list {{
+                display: flex;
+                flex-wrap: wrap;
+                justify-content: space-between;
+            }}
+            .supply-center, .unit {{
+                flex: 0 0 30%;
+                margin-bottom: 10px;
+                padding: 8px;
+                background: #f0f0f0;
+                border-radius: 5px;
+                text-align: center;
+            }}
+            h2 {{
+                border-bottom: 2px solid #eee;
+                padding-bottom: 10px;
+                margin-top: 0;
+            }}
+            .outcome {{
+                background: #e8f5e9;
+                padding: 15px;
+                border-radius: 8px;
+                margin-top: 15px;
+                font-weight: bold;
+                text-align: center;
+            }}
+            .austria {{ border-top: 5px solid #ff5050; }}
+            .england {{ border-top: 5px solid #5050ff; }}
+            .france {{ border-top: 5px solid #50c0ff; }}
+            .germany {{ border-top: 5px solid #808080; }}
+            .italy {{ border-top: 5px solid #50ff50; }}
+            .russia {{ border-top: 5px solid #ffffff; border: 1px solid #ccc; }}
+            .turkey {{ border-top: 5px solid #c0c000; }}
+        </style>
+    </head>
+    <body>
+        <div class="central-info">
+            <h2>Game Information</h2>
+            <div class="game-info">
+                <div class="info-card">
+                    <h3>Game Details</h3>
+                    <p><strong>Game ID:</strong> {game_id}</p>
+                    <p><strong>Phase:</strong> {phase}</p>
+                    <p><strong>Map:</strong> {map_name}</p>
+                    <p><strong>Status:</strong> {status}</p>
+                </div>
+                <div class="info-card">
+                    <h3>Supply Centers</h3>
+                    <div class="supply-centers">
+    """.format(
+        game_id=game_id,
+        phase=phase,
+        map_name=map_name,
+        status="Completed" if is_game_done else "Active"
+    )
+
+    # Add supply center information
+    for power, power_centers in centers.items():
+        html_content += f"""
+                        <div class="supply-center">
+                            <strong>{power}:</strong> {len(power_centers)}
+                        </div>
+        """
+
+    html_content += """
+                    </div>
+                </div>
+            </div>
+    """
+
+    # Add outcome if game is done
+    if is_game_done and outcome:
+        winners = outcome[1:] if len(outcome) > 1 else ["Draw"]
+        html_content += f"""
+            <div class="outcome">
+                <h3>Game Outcome</h3>
+                <p>Winners: {', '.join(winners)}</p>
+            </div>
+        """
+
+    html_content += """
+        </div>
+        <div class="container">
+    """
+
+    # Add each power's information
+    for agent_log in agent_infos:
+        power_name = agent_log["power_name"]
+        power_class = power_name.lower()
+        orders = agent_log.get("orders", [])
+        message_history = agent_log.get("message_history", [])
+        
+        html_content += f"""
+            <div class="power-column {power_class}">
+                <div class="power-name">{power_name}</div>
+                
+                <div class="info-card">
+                    <h3>Units</h3>
+                    <ul>
+        """
+        
+        # Add units information
+        power_units = units.get(power_name, [])
+        for unit in power_units:
+            html_content += f"<li>{unit}</li>"
+        
+        html_content += """
+                    </ul>
+                </div>
+                
+                <div class="message orders">
+                    <div class="role">Final Orders</div>
+                    <ul>
+        """
+        
+        # Add orders
+        for order in orders:
+            html_content += f"<li>{order}</li>"
+        
+        html_content += """
+                    </ul>
+                </div>
+        """
+        
+        # Add message history
+        for message in message_history:
+            if isinstance(message, dict):
+                # Skip system messages or handle differently
+                if message.get("role") == "system":
+                    continue
+                
+                role = message.get("role", "unknown")
+                content = message.get("content", "")
+                
+                role_class = "user" if role == "user" else "assistant"
+                role_display = "Environment" if role == "user" else f"LLM ({power_name})"
+                
+                # Escape HTML characters in content
+                content = content.replace("<", "&lt;").replace(">", "&gt;").replace("\n", "<br>")
+                
+                html_content += f"""
+                <div class="message {role_class}">
+                    <div class="role">{role_display}</div>
+                    <p>{content}</p>
+                </div>
+                """
+            elif isinstance(message, str):
+                # Simple string messages (may be used in some implementations)
+                html_content += f"""
+                <div class="message">
+                    <p>{message}</p>
+                </div>
+                """
+        
+        html_content += """
+            </div>
+        """
+
+    html_content += """
+        </div>
+    </body>
+    </html>
+    """
+    
+    return html_content
+
+def _json_serialize(obj):
+    """
+    A helper function to convert non-JSON-serializable objects
+    (like OrderResult) into strings or dicts.
+    """
+    # Check for the specific object types you know are problematic
+    if obj.__class__.__name__ == "OrderResult":
+        # Return a string representation or a dict
+        return str(obj)
+    
+    # Fallback: attempt to convert anything else to string
+    return str(obj)

src_code_for_reproducibility/markov_games/ipd/Ipd_hard_coded_agents.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+from typing import Any, Tuple
+
+from mllm.markov_games.ipd.ipd_agent import IPDAgent
+from mllm.markov_games.rollout_tree import AgentActLog, ChatTurn
+
+
+@dataclass
+class AlwaysCooperateIPDAgent(IPDAgent):
+	async def act(self, observation) -> Tuple[Any, AgentActLog]:
+		"""
+		Always plays the cooperate action, ignoring observation.
+		Returns the configured cooperate_string so the simulation parses it as "C".
+		"""
+
+		action = self.cooperate_string
+
+		# Log a minimal, structured chat turn for consistency with other agents
+		turn_text = f"Playing cooperate: {action}"
+		self.state.chat_history.append(
+			ChatTurn(
+				agent_id=self.agent_id,
+				role="assistant",
+				content=turn_text,
+				is_state_end=True,
+			)
+		)
+
+		act_log = AgentActLog(
+			chat_turns=[self.state.chat_history[-1]],
+			info=None,
+		)
+
+		# Advance internal counters similar to IPDAgent semantics
+		self.state.chat_counter = len(self.state.chat_history)
+		self.state.round_nb = observation.round_nb
+
+		return action, act_log
+
+
+@dataclass
+class AlwaysDefectIPDAgent(IPDAgent):
+	async def act(self, observation) -> Tuple[Any, AgentActLog]:
+		"""
+		Always plays the defect action, ignoring observation.
+		Returns the configured defect_string so the simulation parses it as "D".
+		"""
+
+		action = self.defect_string
+
+		# Log a minimal, structured chat turn for consistency with other agents
+		turn_text = f"Playing defect: {action}"
+		self.state.chat_history.append(
+			ChatTurn(
+				agent_id=self.agent_id,
+				role="assistant",
+				content=turn_text,
+				is_state_end=True,
+			)
+		)
+
+		act_log = AgentActLog(
+			chat_turns=[self.state.chat_history[-1]],
+			info=None,
+		)
+
+		# Advance internal counters similar to IPDAgent semantics
+		self.state.chat_counter = len(self.state.chat_history)
+		self.state.round_nb = observation.round_nb
+
+		return action, act_log
+

src_code_for_reproducibility/markov_games/ipd/__init__.py

@@ -0,0 +1,7 @@
+from .Ipd_hard_coded_agents import AlwaysCooperateIPDAgent, AlwaysDefectIPDAgent
+
+__all__ = [
+	"AlwaysCooperateIPDAgent",
+	"AlwaysDefectIPDAgent",
+]
+

src_code_for_reproducibility/markov_games/ipd/ipd_agent.py

@@ -0,0 +1,115 @@
+import copy
+import json
+import random
+import re
+from collections.abc import Callable
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from mllm.markov_games.agent import Agent
+from mllm.markov_games.rollout_tree import AgentActLog, ChatTurn
+
+
+@dataclass
+class IPDAgentState:
+    """
+    TOWRITE
+    """
+
+    nb_retries: int
+    round_nb: int
+    chat_counter: int
+    chat_history: List[ChatTurn]
+
+
+@dataclass
+class IPDAgent(Agent):
+    seed: int
+    agent_id: str
+    agent_name: str
+    policy: Callable[[List[Dict]], str]
+    intro_prompt: str  # Introduction prompt explaining the game rules
+    goal_prompt: str  # Prompt explaining the agent's goal
+    strategy_prompt: str  # Prompt suggesting a strategy to the agent
+    max_errors: int  # Maximum number of errors allowed before default action
+    allow_reasoning: bool  # Whether to allow reasoning in the response
+    max_reasoning_chars: int  # Maximum number of characters for reasoning
+    cooperate_string: str  # string parsed as playing cooperate by simulation
+    defect_string: str  # string parsed as playing defect by simulation
+
+    def __post_init__(self):
+        self.state = IPDAgentState(
+            nb_retries=0, round_nb=0, chat_counter=0, chat_history=[]
+        )
+
+    async def act(self, observation) -> Tuple[Any, AgentActLog]:
+        """
+        TOWRITE
+        """
+
+        action = None
+        action_is_ready = False
+        round_nb = observation.round_nb
+
+        # If it's the first round, we need to send the intro prompt
+        if round_nb == 0 and self.state.chat_counter == 0:
+            self.state.chat_history.append(
+                ChatTurn(
+                    agent_id=self.agent_id,
+                    role="user",
+                    content=self.intro_prompt,
+                    is_state_end=True,
+                )
+            )
+
+        # If new round
+        if round_nb > self.state.round_nb:
+            coagent_action = observation.last_coagent_move
+            user_message = f"Last round, the other agent played {coagent_action}."
+            self.state.chat_history.append(
+                ChatTurn(
+                    agent_id=self.agent_id,
+                    role="user",
+                    content=user_message,
+                    is_state_end=True,
+                )
+            )
+
+        # If not new round, try to get valid action from policy
+        output_chat_turn: ChatTurn = await self.policy(
+            state=self.state.chat_history,
+            agent_id=self.agent_id,
+            regex=f"({self.cooperate_string}|{self.defect_string})",
+        )
+        self.state.chat_history.append(output_chat_turn)
+        action = output_chat_turn.content
+
+        agent_step_log = AgentActLog(
+            chat_turns=self.state.chat_history[self.state.chat_counter :], info=None
+        )
+        self.state.chat_counter = len(self.state.chat_history)
+        self.state.round_nb = round_nb
+
+        return action, agent_step_log
+
+    def get_safe_copy(self):
+        """
+        Return a safe copy of the agent.
+        """
+        agent_copy = copy.copy(self)
+        agent_copy.state = copy.deepcopy(self.state)
+        return agent_copy
+
+    def reset(self):
+        self.state = IPDAgentState()
+        raise NotImplementedError
+
+    def render(self):
+        pass
+
+    def close(self):
+        pass
+
+    def get_agent_info(self):
+        pass

src_code_for_reproducibility/markov_games/ipd/ipd_simulation.py

@@ -0,0 +1,162 @@
+import copy
+import random
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from mllm.markov_games.markov_game import Simulation
+from mllm.markov_games.rollout_tree import SimulationStepLog
+from mllm.utils.get_coagent_id import get_coagent_id
+
+
+@dataclass
+class IPDState:
+    """
+    State of the Iterated Prisoner's Dilemma game.
+    """
+
+    round_nb: int = 0
+    done: bool = False
+    last_moves: Dict[str, str] | None = None
+
+
+@dataclass
+class IPDObs:
+    """
+    Observation in Iterated Prisoner's Dilemma game.
+    """
+
+    round_nb: int
+    last_coagent_move: str | None
+
+
+class IPD(Simulation):
+    """
+    Iterated Prisoner's Dilemma simulation following the standard.
+
+    In each round of the game, two agents simultaneously choose to either cooperate (C) or defect (D).
+    The payoffs are as follows:
+    - If both cooperate: Both receive the "reward" (usually 3 points)
+    - If both defect: Both receive the "punishment" (usually 1 point)
+    - If one cooperates and one defects: The defector receives the "temptation" (usually 5 points)
+      and the cooperator receives the "sucker" payoff (usually 0 points)
+
+    The game is played for a specified number of rounds.
+    """
+
+    def __init__(
+        self,
+        agent_ids: List[str],
+        agent_names: List[str],
+        seed: int,
+        rounds_per_game: int,
+        reward: float,  # Both cooperate
+        punishment: float,  # Both defect
+        temptation: float,  # Defector's reward when other cooperates
+        sucker: float,  # Cooperator's reward when other defects
+        cooperate_actions: List[str],
+        defect_actions: List[str],
+    ):
+        self.agent_ids = agent_ids
+        self.agent_names = agent_names
+        self.seed = seed
+        self.rounds_per_game = rounds_per_game
+        self.reward = reward
+        self.punishment = punishment
+        self.temptation = temptation
+        self.sucker = sucker
+        self.cooperate_actions = cooperate_actions
+        self.defect_actions = defect_actions
+        self.state = IPDState()
+
+    def step(self, actions: Dict[str, str]) -> Tuple[bool, SimulationStepLog]:
+        """
+        Take a step in the environment using the provided actions.
+        Here, the observations are just the states of the game.
+
+        Args:
+            actions (dict): A dictionary where keys are agent identifiers and values are actions ('C' or 'D').
+
+        Returns:
+            observations (dict): A dictionary where keys are agent identifiers and values are observations.
+            done (bool): Whether the episode has ended.
+            info (dict): Additional information about the environment.
+        """
+
+        # Calculate rewards using payoff matrix
+        agent0_action = actions[self.agent_ids[0]]
+        agent1_action = actions[self.agent_ids[1]]
+
+        # Normalize actions to standard cooperate/defect/gibberish format
+        def normalize_action(action):
+            if action in self.cooperate_actions:
+                return "C"
+            elif action in self.defect_actions:
+                return "D"
+            else:
+                return "D"
+
+        norm_action0 = normalize_action(agent0_action)
+        norm_action1 = normalize_action(agent1_action)
+
+        payoffs = {
+            ("C", "C"): [self.reward, self.reward],
+            ("C", "D"): [self.sucker, self.temptation],
+            ("D", "C"): [self.temptation, self.sucker],
+            ("D", "D"): [self.punishment, self.punishment],
+        }
+
+        round_rewards = {
+            self.agent_ids[0]: payoffs[(norm_action0, norm_action1)][0],
+            self.agent_ids[1]: payoffs[(norm_action0, norm_action1)][1],
+        }
+
+        # Update game state
+        self.state.round_nb += 1
+        self.state.last_moves = copy.deepcopy(actions)
+        done = self.state.round_nb >= self.rounds_per_game
+        step_log = SimulationStepLog(
+            rewards=round_rewards,
+            info={
+                "actions": {
+                    self.agent_ids[0]: norm_action0,
+                    self.agent_ids[1]: norm_action1,
+                }
+            },
+        )
+
+        return done, step_log
+
+    def get_obs(self):
+        """Returns all agent observations in dict
+        Returns:
+            observations
+        """
+        observations = {}
+        for agent_id in self.agent_ids:
+            observations[agent_id] = self.get_obs_agent(agent_id)
+        return observations
+
+    def get_obs_agent(self, agent_id):
+        """Returns observation for agent_id"""
+        if self.state.last_moves != None:
+            other_id = get_coagent_id(self.agent_ids, agent_id)
+            last_coagent_move = self.state.last_moves[other_id]
+        else:
+            last_coagent_move = None
+        obs = IPDObs(round_nb=self.state.round_nb, last_coagent_move=last_coagent_move)
+        return obs
+
+    def reset(self):
+        """Returns initial observations and states"""
+        self.state = IPDState()
+        return self.get_obs()
+
+    def get_safe_copy(self):
+        """
+        Return a safe copy of the simulation.
+        """
+        simulation_copy = copy.copy(self)
+        simulation_copy.state = copy.deepcopy(self.state)
+        return simulation_copy

src_code_for_reproducibility/markov_games/ipd/ipd_statistics.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import Dict, Callable, List, Tuple
+
+from mllm.markov_games.rollout_tree import SimulationStepLog
+
+
+def avg_reward(sl: SimulationStepLog) -> List[Tuple[str, float]]:
+    for aid in sl.rewards.keys():
+        if "buffer" in str(aid) and "live" not in str(aid):
+            return None
+    # One value per agent at each step
+    rewards_dict = {f"reward-{aid}": float(v) for aid, v in (sl.rewards or {}).items()}
+    return [(key, value) for key, value in rewards_dict.items() if value is not None]
+
+stat_functs: list[Callable[[SimulationStepLog], List[Tuple[str, float]]]] = [
+    avg_reward,
+]

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

src_code_for_reproducibility/markov_games/negotiation/README.md

@@ -0,0 +1,40 @@
+## Negotiation Games: core mechanics and variants
+
+This family of games feature two agents who, in each round, may briefly communicate and then simultaneously propose how to split a fixed resource (most commonly 10 coins). Rewards are the amount kept multiplied by an agent’s per-unit value. The starting speaker alternates deterministically across rounds.
+
+Communication is optional and variant-dependent: some settings encourage rich messaging to share private information, while others remove messaging entirely to focus on allocation behavior.
+
+Proportional splitting is used when the two proposals exceed the available total: allocations are scaled proportionally rather than discarded. This preserves a useful learning signal even when agents over-claim.
+
+### Variants (in increasing difficulty) 
+
+- No‑Press Split 
+  - Single item type (coins)
+  - No communication; agents go straight to making split proposals, with the starting player alternating deterministically.
+  - Motivation: mirrors no‑communication setups (e.g., Advantage Alignment) while keeping the split decision nontrivial.
+  - Deterministic Mode: values are fixed and public: one agent values coins at 10, the other at 1 (alternates each round).
+  - Stochastic Mode: values are random and uncorrelated. 
+
+- Trust-and-Split RPS (TAS-RPS)
+  - Single item type (coins)
+  - Each round, a rock–paper–scissors hand draw creates a strong asymmetry: the winner’s per-coin value is 10, the loser’s is 1.
+  - Each agent initially sees only their own hand and must communicate to coordinate an optimal split.
+  - Motivation: enforce large value disparity so one’s own value reveals little about the other’s (avoiding ceiling effects) and incentivize meaningful communication.
+ 
+- Trust-and-Split (TAS)
+  - Single item type (coins); each round, each agent’s per-coin value is independently sampled in a broad range (e.g., 1–20).
+  - Each agent observes only their own value; they may use short messages to share and negotiate.
+  - Motivation: a simple blend that tests whether agents learn to exchange private information and coordinate proportional, value-aware splits.
+
+- Deal-or-No-Deal (DOND)
+  - Introduced in [Deal or No Deal? End-to-End Learning for Negotiation Dialogues](https://arxiv.org/pdf/1706.05125)
+  - Multiple item types (typically "books", "hats" and "balls") with limited stocks; each agent has its own per-type values.
+  - A deal pays out only if both proposals exactly agree and respect the stock; otherwise no deal (zero reward) that round.
+  - Motivation: a known benchmark closer to real-world bargaining, where both parties must explicitly agree.
+
+
+
+
+
+
+

src_code_for_reproducibility/markov_games/negotiation/dond_agent.py

@@ -0,0 +1,61 @@
+import copy
+import re
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple
+
+from mllm.markov_games.agent import Agent
+from mllm.markov_games.rollout_tree import AgentActLog, ChatTurn
+from mllm.markov_games.negotiation.dond_simulation import (
+    DealNoDealObs,
+)
+from mllm.markov_games.negotiation.nego_simulation import Split
+from mllm.markov_games.negotiation.nego_agent import NegotiationAgent, NegotiationAgentState
+
+class DealNoDealAgent(NegotiationAgent):
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.intro_prompt = (
+                "You are {agent_id}. You are playing an iterated game. "
+                "At each round, you and other agent will try to distribute among yourselves items of types {item_types}. "
+                "You only know how much you value each item type, but not the other agent's values. "
+                "You can communicate with the other agent by sending up to {quota_messages_per_agent_per_round} short messages per round. "
+                "Each round, after exchanging messages, you and the other agent will submit a private proposal. "
+                "A deal is accepted only if both proposals match exactly and are within stock; otherwise no deal (0 points for both at that round). "
+                "The values of the items of the other agent at the previous round are revealed to you after each round. "
+                "Your goal is: {goal}."
+            )
+        self.new_round_prompt = ("New round {round_nb}. Items: {stock}. Your values: {values}. ")
+        self.last_round_prompt = ("Last round, other agent's values: {previous_values_coagent}. ")
+        self.send_split_prompt = ("Respond with <split>...</split> where you propose how many items of each type you want to keep.")
+        
+    def get_message_regex(self, observation: DealNoDealObs) -> str:
+        return r"<message>[\s\S]{0,400}</message>"
+    
+    def get_split_regex(self, observation: DealNoDealObs) -> str:
+        parts = []
+        for t in observation.item_types:
+            s = int(observation.quantities.get(t, 0))
+            allowed = "|".join(str(k) for k in range(0, s + 1))
+            rng = f"({allowed})"
+            parts.append(fr"<{t}>{rng}</{t}>")
+        items_block = "".join(parts)
+        return fr"(<split>{items_block}</split>)"
+    
+    def get_split_action(self, policy_output: str, observation: DealNoDealObs) -> Split:
+        import re as _re
+        allocations: Dict[str, int] = {}
+        for t in observation.item_types:
+            m = _re.search(fr"<{t}>([0-9]+)</{t}>", policy_output)
+            if m:
+                allocations[t] = int(m.group(1))
+            else:
+                allocations[t] = 0
+        return Split(items_given_to_self=allocations)
+ 
+
+

src_code_for_reproducibility/markov_games/negotiation/dond_simulation.py

@@ -0,0 +1,153 @@
+import copy
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple
+
+from numpy.random import default_rng
+
+from mllm.markov_games.rollout_tree import SimulationStepLog
+from mllm.markov_games.negotiation.nego_simulation import Split, NegotiationState, NegotiationObs, NegotiationSimulation
+from mllm.utils.get_coagent_id import get_coagent_id
+
+
+AgentId = str
+
+
+@dataclass
+class DealNoDealState(NegotiationState):
+    item_types: List[str]
+    values: Dict[AgentId, Dict[str, int]]
+
+@dataclass
+class DealNoDealObs(NegotiationObs):
+    my_values: Dict[str, int]
+    item_types: List[str]
+    previous_values_coagent: Dict[str, int] | None
+
+
+def random_partition_integer(rng, total: int, parts: int) -> List[int]:
+    if parts <= 0:
+        return []
+    if total <= 0:
+        return [0 for _ in range(parts)]
+    cuts = sorted(rng.integers(0, total + 1, size=parts - 1).tolist())
+    vals = []
+    prev = 0
+    for c in cuts + [total]:
+        vals.append(c - prev)
+        prev = c
+    return vals
+
+class DealNoDealSimulation(NegotiationSimulation):
+
+    def __init__(
+        self,
+        item_types: List[str] = ["books", "hats", "balls"],
+        *args,
+        **kwargs,
+    ):
+        super().__init__(item_types=item_types, *args, **kwargs)
+        self.reset()
+
+    def _other(self, agent_id: AgentId) -> AgentId:
+        return get_coagent_id(self.agent_ids, agent_id)
+
+    def _sample_stock(self) -> Dict[str, int]:
+        # total items between 5 and 7
+        total_items = int(self.rng.integers(5, 8))
+        # nonnegative per-type counts summing to total_items
+        parts = random_partition_integer(self.rng, total_items, len(self.item_types))
+        # allow zeros per type
+        return {t: int(c) for t, c in zip(self.item_types, parts)}
+
+    def _sample_values_pair(self) -> Dict[AgentId, Dict[str, int]]:
+        # Each agent has integer non-negative values that sum to 10
+        # Each item type valued by at least one agent
+        # Some item type valued by both agents
+        while True:
+            vals_a = random_partition_integer(self.rng, 10, len(self.item_types))
+            vals_b = random_partition_integer(self.rng, 10, len(self.item_types))
+            a = {t: int(v) for t, v in zip(self.item_types, vals_a)}
+            b = {t: int(v) for t, v in zip(self.item_types, vals_b)}
+            # each item valued by at least one
+            ok1 = all((a[t] > 0) or (b[t] > 0) for t in self.item_types)
+            # some item valued by both
+            ok2 = any((a[t] > 0) and (b[t] > 0) for t in self.item_types)
+            if ok1 and ok2:
+                return {self.agent_ids[0]: a, self.agent_ids[1]: b}
+
+    def _is_valid_allocation(self, allocation: Dict[str, int], stock: Dict[str, int]) -> bool:
+        for t in self.item_types:
+            v = allocation.get(t)
+            if v is None:
+                return False
+            if not isinstance(v, int):
+                return False
+            if v < 0 or v > int(stock.get(t, 0)):
+                return False
+        return True
+    
+    def set_new_round_of_variant(self):
+        # Keep same values, resample stock
+        self.state.quantities = self._sample_stock()
+
+    def get_info_of_variant(self, state: NegotiationState, actions: Dict[AgentId, Any]) -> Dict[str, Any]:
+        return {
+            "quantities": copy.deepcopy(state.quantities),
+            "values": copy.deepcopy(state.values),
+            'splits': copy.deepcopy(state.splits),
+        }
+
+    def get_rewards(self, splits: Dict[AgentId, Split]) -> Dict[AgentId, float]:
+        """
+        Returns the rewards for each agent.
+        """
+        split_a = splits[self.agent_ids[0]].items_given_to_self
+        split_b = splits[self.agent_ids[1]].items_given_to_self
+        rewards = {self.agent_ids[0]: 0, self.agent_ids[1]: 0}
+        for t in self.item_types:
+            # If not complementary, return 0! 
+            if not split_a[t] + split_b[t] == self.state.quantities[t]:
+                return {self.agent_ids[0]: 0, self.agent_ids[1]: 0}
+            rewards[self.agent_ids[0]] += split_a[t] * self.state.values[self.agent_ids[0]][t]
+            rewards[self.agent_ids[1]] += split_b[t] * self.state.values[self.agent_ids[1]][t]
+        return rewards
+
+    def get_obs(self):
+        return {agent_id: self.get_obs_agent(agent_id) for agent_id in self.agent_ids}
+
+    def get_obs_agent(self, agent_id):
+        other_id = self._other(agent_id)
+        obs = DealNoDealObs(
+            round_nb=self.state.round_nb,
+            last_message=self.state.last_message,
+            current_agent=self.state.current_agent,
+            quantities=copy.deepcopy(self.state.quantities),
+            value=0.0,  # unused in DOND
+            other_agent_split=None,  # not meaningful until split
+            split_phase=self.state.split_phase,
+            quota_messages_per_agent_per_round=self.quota_messages_per_agent_per_round,
+            my_values=copy.deepcopy(self.state.values[agent_id]),
+            item_types=list(self.item_types),
+            previous_values_coagent=copy.deepcopy(self.state.values.get(other_id, {})),
+        )
+        return obs
+
+    def reset(self):
+        start_agent = self.agent_ids[self._starting_agent_index]
+        stock = self._sample_stock()
+        values = self._sample_values_pair()
+        self.state = DealNoDealState(
+            round_nb=0,
+            last_message="",
+            current_agent=start_agent,
+            quantities=stock,
+            values=values,
+            previous_values=None,
+            splits={aid: None for aid in self.agent_ids},
+            nb_messages_sent={aid: 0 for aid in self.agent_ids},
+            split_phase=False,
+            item_types=list(self.item_types),
+        )
+        return self.get_obs()
+
+

src_code_for_reproducibility/markov_games/negotiation/nego_agent.py

@@ -0,0 +1,242 @@
+import copy
+from abc import abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+
+from mllm.markov_games.agent import Agent
+from mllm.markov_games.negotiation.nego_simulation import Message, NegotiationObs, Split
+from mllm.markov_games.rollout_tree import AgentActLog, ChatTurn
+
+
+@dataclass
+class NegotiationAgentState:
+    round_nb: int
+    nb_messages_sent_this_round: int
+    chat_counter: int
+    chat_history: List[ChatTurn]
+
+
+class NegotiationAgent(Agent):
+    def __init__(
+        self,
+        seed: int,
+        agent_id: str,
+        agent_name: str,
+        policy: Callable[[List[Dict]], str],
+        goal: str,
+        exploration_prompts: List[str] = [],
+        exploration_prompt_probs: List[float] = [],
+    ):
+        self.seed = seed
+        self.agent_id = agent_id
+        self.agent_name = agent_name
+        self.policy = policy
+        self.goal = goal
+        self.exploration_prompts_toggled = len(exploration_prompts) > 0
+        if self.exploration_prompts_toggled:
+            exploration_prompts = copy.deepcopy(exploration_prompts)
+            exploration_prompts.append(None)
+            self.exploration_prompts = exploration_prompts
+            self.exploration_prompt_probs = np.array(exploration_prompt_probs)
+            assert self.exploration_prompt_probs.sum() <= 1
+            assert np.all(self.exploration_prompt_probs >= 0)
+            self.exploration_prompt_probs = np.append(
+                self.exploration_prompt_probs, 1 - self.exploration_prompt_probs.sum()
+            )
+        self.state = NegotiationAgentState(
+            round_nb=0, nb_messages_sent_this_round=0, chat_counter=0, chat_history=[]
+        )
+
+        # Implemented in variants
+        self.intro_prompt = ""
+        self.new_round_prompt = ""
+        self.last_round_prompt = ""
+        self.send_split_prompt = ""
+        self.wait_for_message_prompt = ""
+        self.last_message_prompt = ""
+        self.send_message_prompt = ""
+
+    @abstractmethod
+    def get_message_regex(self, observation: NegotiationObs) -> str:
+        pass
+
+    @abstractmethod
+    def get_split_regex(self, observation: NegotiationObs) -> str:
+        pass
+
+    @abstractmethod
+    def get_split_action(
+        self, policy_output: str, observation: NegotiationObs
+    ) -> Split:
+        pass
+
+    async def act(self, observation: NegotiationObs) -> Tuple[Any, AgentActLog]:
+        def dict_to_str(d: dict) -> str:
+            return ", ".join(f"{v} {k}" for k, v in d.items())
+
+        def dict_to_eq_str(d: dict) -> str:
+            return ", ".join(f"{k}={v}" for k, v in d.items())
+
+        is_our_turn = observation.current_agent == self.agent_id
+        action: Any = None
+        round_nb = observation.round_nb
+
+        prompt_parts: List[str] = []
+        obs_ctx = vars(observation)
+        obs_ctx_formmated = obs_ctx.copy()
+        for key in obs_ctx_formmated:
+            if isinstance(obs_ctx_formmated[key], dict) and "value" not in key:
+                obs_ctx_formmated[key] = dict_to_str(obs_ctx_formmated[key])
+            elif isinstance(obs_ctx_formmated[key], dict) and "value" in key:
+                obs_ctx_formmated[key] = dict_to_eq_str(obs_ctx_formmated[key])
+
+        #######################################
+        # build user prompt
+        #######################################
+
+        # First-ever call
+        is_intro = round_nb == 0 and self.state.chat_counter == 0
+        if is_intro:
+            prompt_parts.append(
+                self.intro_prompt.format(
+                    goal=self.goal, agent=self.agent_name, **obs_ctx_formmated
+                )
+            )
+
+        # New round
+        is_new_round = round_nb > self.state.round_nb
+        if is_new_round or is_intro:
+            self.state.nb_messages_sent_this_round = 0
+            if not is_intro:
+                prompt_parts.append(self.last_round_prompt.format(**obs_ctx_formmated))
+            prompt_parts.append(self.new_round_prompt.format(**obs_ctx_formmated))
+            if self.exploration_prompts_toggled:
+                exploration_prompt = self.exploration_prompts[
+                    np.random.choice(
+                        len(self.exploration_prompts), p=self.exploration_prompt_probs
+                    )
+                ]
+                if exploration_prompt is not None:
+                    prompt_parts.append(exploration_prompt)
+            self.state.round_nb = round_nb
+
+        # Wait for message
+        if not is_our_turn and not observation.split_phase:
+            prompt_parts.append(
+                self.wait_for_message_prompt.format(**obs_ctx_formmated)
+            )
+
+        # Get last message
+        if is_our_turn and not is_new_round and not is_intro:
+            prompt_parts.append(self.last_message_prompt.format(**obs_ctx_formmated))
+
+        # Prompt to send message
+        must_send_message = not observation.split_phase and is_our_turn
+        if must_send_message:
+            prompt_parts.append(self.send_message_prompt.format(**obs_ctx_formmated))
+
+        # Prompt to give split
+        must_send_split = not must_send_message and observation.split_phase
+        if must_send_split:
+            var_names = ["x", "y", "z", "w"]  # Extend as needed
+            items_str = ", ".join(
+                [
+                    f"{var_names[i]} {item}"
+                    for i, item in enumerate(obs_ctx["quantities"].keys())
+                ]
+            )
+            ranges_str = ", ".join(
+                [
+                    f"{var_names[i]}: 0-{obs_ctx['quantities'][item]} (integer)"
+                    for i, item in enumerate(obs_ctx["quantities"].keys())
+                ]
+            )
+            proposal_style = f"Proposal: {items_str} where {ranges_str}."
+            proposal_style2 = (
+                f"<items_to_self> {items_str} </items_to_self> where {ranges_str}."
+            )
+            prompt_parts.append(
+                self.send_split_prompt.format(
+                    proposal_style=proposal_style,
+                    proposal_style2=proposal_style2,
+                    **obs_ctx_formmated,
+                )
+            )
+
+        # Append one ChatTurn with is_state_end=True
+        user_prompt = "\n".join(prompt_parts)
+        self.state.chat_history.append(
+            ChatTurn(
+                agent_id=self.agent_id,
+                role="user",
+                content=user_prompt,
+                is_state_end=True,
+            )
+        )
+
+        #######################################
+        # Get policy action
+        #######################################
+
+        # Query policy for the appropriate format
+        if must_send_message:
+            return_regex = self.get_message_regex(observation)
+            policy_output = await self.policy(
+                state=self.state.chat_history,
+                agent_id=self.agent_id,
+                regex=return_regex,
+            )
+            self.state.chat_history.append(
+                ChatTurn(
+                    agent_id=self.agent_id,
+                    role="assistant",
+                    content=policy_output.content,
+                    reasoning_content=policy_output.reasoning_content,
+                    log_probs=policy_output.log_probs,
+                    out_token_ids=policy_output.out_token_ids,
+                    is_state_end=False,
+                )
+            )
+            action = Message(message=policy_output.content)
+            self.state.nb_messages_sent_this_round += 1
+
+        elif must_send_split:
+            return_regex = self.get_split_regex(observation)
+            policy_output = await self.policy(
+                state=self.state.chat_history,
+                agent_id=self.agent_id,
+                regex=return_regex,
+            )
+            self.state.chat_history.append(
+                ChatTurn(
+                    agent_id=self.agent_id,
+                    role="assistant",
+                    content=policy_output.content,
+                    reasoning_content=policy_output.reasoning_content,
+                    log_probs=policy_output.log_probs,
+                    out_token_ids=policy_output.out_token_ids,
+                    is_state_end=False,
+                )
+            )
+            action = self.get_split_action(policy_output.content, observation)
+        else:
+            action = None
+
+        agent_step_log = AgentActLog(
+            chat_turns=self.state.chat_history[self.state.chat_counter :], info=None
+        )
+        self.state.chat_counter = len(self.state.chat_history)
+        return action, agent_step_log
+
+    def get_safe_copy(self):
+        agent_copy = copy.copy(self)
+        agent_copy.state = copy.deepcopy(self.state)
+        return agent_copy
+
+    def reset(self):
+        self.state = NegotiationAgentState(
+            round_nb=0, nb_messages_sent_this_round=0, chat_counter=0, chat_history=[]
+        )

src_code_for_reproducibility/markov_games/negotiation/nego_hard_coded_policies.py

@@ -0,0 +1,64 @@
+import asyncio
+from typing import Optional
+from mllm.markov_games.negotiation.nego_agent import NegotiationAgent
+from mllm.markov_games.negotiation.no_press_nego_agent import NoPressAgent
+from mllm.markov_games.negotiation.no_press_nego_simulation import NoPressObs
+from mllm.markov_games.rollout_tree import AgentActLog, ChatTurn
+from mllm.markov_games.negotiation.nego_simulation import Split
+from typing import Any, Tuple
+
+class HardCodedNegoWelfareMaximizingPolicy(NoPressAgent):
+    async def act(self, observation: NoPressObs) -> Tuple[Any, AgentActLog]:
+        """
+        Policy that gives all of the items to the agent who values them more.
+        If the items are equally valued, give them to the agent who values them more.
+        """
+        quantities = observation.quantities
+        my_values = observation.value
+        other_values = observation.other_value
+
+        items_given_to_self = {}
+        for item, qty in quantities.items():
+            my_v = float(my_values.get(item, 0))
+            other_v = float(other_values.get(item, 0))
+            if my_v == other_v:
+                items_given_to_self[item] = int(qty) / 2
+            else:
+                items_given_to_self[item] = int(qty if my_v > other_v else 0)
+
+        action = Split(items_given_to_self=items_given_to_self)
+        act_log = AgentActLog(
+            chat_turns=[
+                ChatTurn(
+                    agent_id=self.agent_id,
+                    role="assistant",
+                    content="Using welfare-maximizing split (all to higher-value agent).",
+                    is_state_end=True,
+                )
+            ],
+            info=None,
+        )
+        return action, act_log
+
+class HardCodedNegoGreedyPolicy(NoPressAgent):
+    async def act(self, observation: NoPressObs) -> Tuple[Any, AgentActLog]:
+        """
+        Always gives itself all of the items.
+        """
+        quantities = observation.quantities
+        items_given_to_self = {item: int(qty) for item, qty in quantities.items()}
+
+        action = Split(items_given_to_self=items_given_to_self)
+        act_log = AgentActLog(
+            chat_turns=[
+                ChatTurn(
+                    agent_id=self.agent_id,
+                    role="assistant",
+                    content="Using greedy split (keep all items).",
+                    is_state_end=True,
+                )
+            ],
+            info=None,
+        )
+        return action, act_log
+

src_code_for_reproducibility/markov_games/negotiation/nego_simulation.py

@@ -0,0 +1,241 @@
+"""
+Negotiation simulation environment
+other agent is set at the start of every round. Even though current agent changes over message turns in a round.
+"""
+import copy
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple
+
+from numpy.random import default_rng
+
+from mllm.markov_games.rollout_tree import SimulationStepLog
+from mllm.markov_games.simulation import Simulation
+from mllm.utils.get_coagent_id import get_coagent_id
+
+AgentId = str
+
+
+@dataclass
+class Split:
+    items_given_to_self: Dict[str, int]
+
+
+@dataclass
+class Message:
+    message: str
+
+
+@dataclass  # gets extended by variants
+class NegotiationState:
+    round_nb: int
+    last_message: str
+    current_agent: AgentId
+    quantities: Dict[str, int]
+    values: Dict[AgentId, Dict[str, float]]
+    splits: Dict[AgentId, Split | None]
+    nb_messages_sent: Dict[AgentId, int]
+    previous_values: Dict[AgentId, Dict[str, float]] | None
+    previous_splits: Dict[AgentId, Dict[str, int] | None] | None
+    previous_points: Dict[AgentId, float] | None
+    previous_quantities: Dict[str, int] | None
+    split_phase: bool
+
+
+@dataclass  # gets extended by variants
+class NegotiationObs:
+    round_nb: int
+    last_message: str
+    quota_messages_per_agent_per_round: int
+    current_agent: AgentId
+    other_agent: str
+    quantities: Dict[str, int]
+    item_types: List[str]
+    value: Dict[str, int]
+    split_phase: bool
+    last_split_agent: Dict[str, int] | None
+    last_value_agent: Dict[str, int] | None
+    last_points_agent: float | None
+    last_split_coagent: Dict[str, int] | None
+    last_value_coagent: Dict[str, int] | None
+    last_points_coagent: float | None
+    last_quantities: Dict[str, int] | None
+
+
+def compute_tas_style_rewards(
+    agent_ids: List[AgentId],
+    values: Dict[AgentId, float],
+    splits: Dict[AgentId, Split],
+    quantities: Dict[str, int],
+) -> Dict[AgentId, float]:
+    """
+    TAS-like reward computation: if sum of proposed coins exceeds max_coins,
+    allocate proportionally. Otherwise, use proposed amounts directly.
+    Rewards are quantity_kept * per-coin value for each agent.
+    """
+    a0, a1 = agent_ids[0], agent_ids[1]
+    r0, r1 = 0.0, 0.0
+
+    for item in quantities:
+        max_item = quantities[item]
+        item_to_self_0 = int(
+            (splits[a0].items_given_to_self.get(item, 0))
+            if splits[a0] is not None
+            else 0
+        )
+        item_to_self_1 = int(
+            (splits[a1].items_given_to_self.get(item, 0))
+            if splits[a1] is not None
+            else 0
+        )
+        denom = max(int(max_item), item_to_self_0 + item_to_self_1)
+        q0 = float(max_item) * float(item_to_self_0) / float(denom)
+        q1 = float(max_item) * float(item_to_self_1) / float(denom)
+        if type(values[a0]) is not dict:
+            r0 += q0 * float(values[a0])
+            r1 += q1 * float(values[a1])
+        else:
+            r0 += q0 * float(values[a0][item])
+            r1 += q1 * float(values[a1][item])
+    return {a0: r0, a1: r1}
+
+
+class NegotiationSimulation(Simulation):
+    def __init__(
+        self,
+        agent_ids: List[AgentId],
+        agent_names: List[str],
+        seed: int,
+        nb_of_rounds: int,
+        quota_messages_per_agent_per_round: int,
+        item_types: List[str] | None = None,
+    ):
+        self.seed = seed
+        self.rng = default_rng(self.seed)
+        self.agent_ids = list(agent_ids)
+        self.agent_names = agent_names
+        self.agent_id_to_name = {
+            agent_id: agent_name for agent_id, agent_name in zip(agent_ids, agent_names)
+        }
+        self.nb_of_rounds = int(nb_of_rounds)
+        self.quota_messages_per_agent_per_round = int(
+            quota_messages_per_agent_per_round
+        )
+        if item_types is not None:
+            self.item_types = [item.lower() for item in item_types]
+        else:
+            self.item_types = ["coins"]
+        self.state: NegotiationState | None = None
+        self._starting_agent_index = self.rng.choice([0, 1])
+        self.reset()
+
+    def _other(self, agent_id: AgentId) -> AgentId:
+        return get_coagent_id(self.agent_ids, agent_id)
+
+    @abstractmethod
+    def set_new_round_of_variant(self):
+        pass
+
+    @abstractmethod
+    def get_info_of_variant(
+        self, state: NegotiationState, actions: Dict[AgentId, Any]
+    ) -> Dict[str, Any]:
+        pass
+
+    def step(self, actions: Any) -> Tuple[bool, SimulationStepLog]:
+        """
+        Returns terminated, step_log
+        """
+        assert self.state is not None
+        current_agent = self.state.current_agent
+        a0, a1 = self.agent_ids[0], self.agent_ids[1]
+        action = actions.get(current_agent)
+
+        # Split phase: require both splits in the same timestep
+        if self.state.split_phase:
+            action_a0 = actions.get(a0)
+            action_a1 = actions.get(a1)
+            have_both_splits = isinstance(action_a0, Split) and isinstance(
+                action_a1, Split
+            )
+            if not have_both_splits:
+                rewards = {agent_id: 0.0 for agent_id in self.agent_ids}
+                return False, SimulationStepLog(
+                    rewards=rewards, info={"type": "waiting_for_splits"}
+                )
+
+            # Record splits
+            self.state.splits[a0] = action_a0
+            self.state.splits[a1] = action_a1
+
+            # Compute rewards and end round
+            rewards = self.get_rewards(self.state.splits)
+
+            # Info
+            info = self.get_info_of_variant(self.state, actions)
+
+            # Prepare next round
+            # Alternate starting agent
+            self.state.round_nb += 1
+            self._starting_agent_index = 1 - self._starting_agent_index
+            self.state.current_agent = self.agent_ids[self._starting_agent_index]
+            self.state.previous_values = copy.deepcopy(self.state.values)
+            self.state.previous_splits = copy.deepcopy(self.state.splits)
+            self.state.previous_quantities = copy.deepcopy(self.state.quantities)
+            self.state.previous_points = copy.deepcopy(rewards)
+            self.state.last_message = ""
+            self.set_new_round_of_variant()  # variant specific
+            self.state.splits = {agent_id: None for agent_id in self.agent_ids}
+            self.state.nb_messages_sent = {agent_id: 0 for agent_id in self.agent_ids}
+            is_last_timestep_in_round = True
+            done = self.state.round_nb >= self.nb_of_rounds
+
+        # Message phase
+        elif isinstance(action, Message):
+            self.state.last_message = action.message
+            self.state.nb_messages_sent[current_agent] += 1
+
+            # Move turn to other agent
+            self.state.current_agent = self._other(current_agent)
+
+            # If both agents have reached their message quota, enter split phase
+            if all(
+                self.state.nb_messages_sent[agent_id]
+                >= self.quota_messages_per_agent_per_round
+                for agent_id in self.agent_ids
+            ):
+                self.state.split_phase = True
+            is_last_timestep_in_round = False
+            done = False
+            rewards = {agent_id: 0.0 for agent_id in self.agent_ids}
+            info = {"type": "message"}
+
+        info[
+            "is_last_timestep_in_round"
+        ] = is_last_timestep_in_round  # Used later to group round timesteps if needed
+        return done, SimulationStepLog(rewards=rewards, info=info)
+
+    def get_obs(self):
+        """Returns all agent observations in dict"""
+        return {agent_id: self.get_obs_agent(agent_id) for agent_id in self.agent_ids}
+
+    @abstractmethod
+    def get_rewards(self, splits: Dict[AgentId, Split]) -> Dict[AgentId, float]:
+        pass
+
+    @abstractmethod
+    def get_obs_agent(self, agent_id):
+        pass
+
+    def get_state(self):
+        return self.state
+
+    def get_safe_copy(self):
+        """Return a safe copy of the simulation."""
+        simulation_copy = copy.copy(self)
+        simulation_copy.state = copy.deepcopy(self.state)
+        return simulation_copy
+
+    @abstractmethod
+    def reset(self) -> dict[AgentId, NegotiationObs]:
+        pass

src_code_for_reproducibility/markov_games/negotiation/no_press_nego_simulation.py

@@ -0,0 +1,168 @@
+import copy
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Dict, List, Literal, Tuple
+
+from mllm.markov_games.negotiation.nego_simulation import (
+    NegotiationObs,
+    NegotiationSimulation,
+    NegotiationState,
+    Split,
+    compute_tas_style_rewards,
+)
+
+AgentId = str
+
+
+@dataclass
+class NoPressState(NegotiationState):
+    pass
+
+
+@dataclass
+class NoPressObs(NegotiationObs):
+    other_value: Dict[str, float]
+
+
+class NoPressSimulation(NegotiationSimulation):
+    def __init__(
+        self,
+        game_type: Literal["10-1-exclusive", "10-1-ties", "1-to-20"] = "1-to-20",
+        same_round_value: bool = True,
+        atleast_one_conflict: bool = False,
+        *args,
+        **kwargs,
+    ):
+        self.game_type = game_type
+        self.same_round_value = same_round_value
+        self.atleast_one_conflict = atleast_one_conflict
+        super().__init__(*args, **kwargs)
+
+    def _sample_values(self) -> Dict[AgentId, dict]:
+        values = defaultdict(dict)
+        if self.state is None:
+            item_types = self.item_types
+        else:
+            item_types = list(self.state.quantities.keys())
+        while True:
+            for item in item_types:
+                if self.game_type == "10-1-exclusive":
+                    v = int(self.rng.choice([1, 10]))
+                    values[self.agent_ids[0]][item] = v
+                    values[self.agent_ids[1]][item] = 10 if v == 1 else 1
+                elif self.game_type == "10-1-ties":
+                    for aid in self.agent_ids:
+                        values[aid][item] = int(self.rng.choice([1, 10]))
+                elif self.game_type == "1-to-20":
+                    for aid in self.agent_ids:
+                        values[aid][item] = int(self.rng.integers(1, 21))
+            if self.atleast_one_conflict:
+                has_conflict = False
+                for item in item_types:
+                    agent_values_for_item = [
+                        values[aid][item] for aid in self.agent_ids
+                    ]
+                    if len(set(agent_values_for_item)) > 1:
+                        has_conflict = True
+                        break
+                if not has_conflict:
+                    continue
+            agent_values = [sum(v.values()) for v in values.values()]
+            if len(set(agent_values)) == 1 or not self.same_round_value:
+                break
+        return values
+
+    def _sample_quantities(self) -> Dict[str, int]:
+        return {item.lower(): 10 for item in self.item_types}
+
+    def set_new_round_of_variant(self):
+        self.state.quantities = self._sample_quantities()
+        self.state.values = self._sample_values()
+        self.state.split_phase = True
+
+    def get_info_of_variant(
+        self, state: NegotiationState, actions: Dict[AgentId, Any]
+    ) -> Dict[str, Any]:
+        return {
+            "quantities": copy.deepcopy(state.quantities),
+            "values": copy.deepcopy(state.values),
+            "splits": copy.deepcopy(state.splits),
+        }
+
+    def get_rewards(self, splits: Dict[AgentId, Split]) -> Dict[AgentId, float]:
+        return compute_tas_style_rewards(
+            self.agent_ids, self.state.values, splits, self.state.quantities
+        )
+
+    def get_obs(self):
+        return {agent_id: self.get_obs_agent(agent_id) for agent_id in self.agent_ids}
+
+    def get_obs_agent(self, agent_id):
+        other_id = self._other(agent_id)
+        last_value_coagent = (
+            None
+            if self.state.previous_values is None
+            else self.state.previous_values.get(other_id)
+        )
+        last_points_coagent = (
+            None
+            if self.state.previous_points is None
+            else round(self.state.previous_points.get(other_id), 1)
+        )
+        last_value_agent = (
+            None
+            if self.state.previous_values is None
+            else self.state.previous_values.get(agent_id)
+        )
+        last_points_agent = (
+            None
+            if self.state.previous_points is None
+            else round(self.state.previous_points.get(agent_id), 1)
+        )
+        last_split_coagent = None
+        last_split_agent = None
+        if self.state.previous_splits is not None:
+            last_split_coagent = self.state.previous_splits[
+                other_id
+            ].items_given_to_self
+            last_split_agent = self.state.previous_splits[agent_id].items_given_to_self
+        obs = NoPressObs(
+            round_nb=self.state.round_nb,
+            last_message="",
+            quota_messages_per_agent_per_round=self.quota_messages_per_agent_per_round,
+            current_agent=self.state.current_agent,
+            other_agent=self.agent_id_to_name[other_id],
+            quantities=self.state.quantities,
+            item_types=self.item_types,
+            value=self.state.values[agent_id],
+            split_phase=self.state.split_phase,
+            last_split_agent=last_split_agent,
+            last_value_agent=last_value_agent,
+            last_points_agent=last_points_agent,
+            last_split_coagent=last_split_coagent,
+            last_value_coagent=last_value_coagent,
+            last_points_coagent=last_points_coagent,
+            other_value=self.state.values[other_id],
+            last_quantities=self.state.previous_quantities,
+        )
+        return obs
+
+    def reset(self):
+        start_agent = self.agent_ids[self._starting_agent_index]
+        quantities = self._sample_quantities()
+        values = self._sample_values()
+        self.state = NoPressState(
+            round_nb=0,
+            last_message="",
+            current_agent=start_agent,
+            quantities=quantities,
+            values=values,
+            previous_values=None,
+            splits={aid: None for aid in self.agent_ids},
+            nb_messages_sent={aid: 0 for aid in self.agent_ids},
+            split_phase=True,
+            previous_splits=None,
+            previous_points=None,
+            previous_quantities=None,
+        )
+        return self.get_obs()

src_code_for_reproducibility/markov_games/negotiation/tas_agent.py

@@ -0,0 +1,108 @@
+from mllm.markov_games.negotiation.nego_agent import NegotiationAgent
+from mllm.markov_games.negotiation.nego_simulation import Split
+from mllm.markov_games.negotiation.tas_simulation import TrustAndSplitObs
+
+
+class TrustAndSplitAgent(NegotiationAgent):
+    def __init__(self, num_message_chars, *args, **kwargs):
+        self.num_message_chars = num_message_chars
+        super().__init__(*args, **kwargs)
+        self.intro_prompt = (
+            "Welcome to an iterated game. You are {agent}. The other agent is {other_agent}.\n"
+            "Setup:\n"
+            "1. The game has multiple independent rounds.\n"
+            "2. In each round, there are multiple items to split between the two agents.\n"
+            "3. Both agents are assigned a per-item value between 1 and 20 (inclusive) in each round.\n"
+            "4. You can only observe your own per-item values.\n"
+            "5. Because assignments are random, both agents are equally likely to have same expected per-item value.\n"
+            "\n"
+            "Protocol:\n"
+            "1. At the start of the round, one agent begins the conversation. The starting role alternates each round.\n"
+            "2. Agents exchange a short chat ({quota_messages_per_agent_per_round} messages per round per agent) to negotiate how to split the item.\n"
+            "   - Use this chat to communicate your private per-item value to make informed proposals.\n"
+            "3. After the chat, both agents simultaneously propose the amount of each item they will keep.\n"
+            "4. If the total sum of proposals is less than or equal to the item quantity, both agents receive their proposed amounts.\n"
+            "5. If the total sum of proposals exceeds the item quantity, they are allocated proportionally.\n"
+            "6. Your points for the round = (amount you receive per item) x (your per-item value for that round), added across all items.\n"
+            "7. Points are accumulated across rounds.\n"
+            "Your goal: {goal}\n"
+        )
+        self.new_round_prompt = (
+            "A New Round Begins\n"
+            "The items to split are {quantities}.\n"
+            "Your per-item values are {value}."
+        )
+        self.last_round_prompt = (
+            "Last Round Summary:\n"
+            "   - Items to split: {last_quantities}\n"
+            "   - Your per-item values: {last_value_agent}\n"
+            "   - {other_agent}'s per-item values: {last_value_coagent}\n"
+            "   - You proposed: {last_split_agent}\n"
+            "   - You earned: {last_points_agent} points\n"
+            "   - {other_agent} proposed: {last_split_coagent}\n"
+            "   - {other_agent} earned: {last_points_coagent} points\n"
+            "   - Round Complete.\n"
+        )
+        self.send_split_prompt = (
+            "Message quota is finished for this round.\n"
+            "{other_agent} has finalized their proposal.\n"
+            "Submit your finalization now\n"
+            "Respond with {proposal_style2}"
+        )
+        # self.wait_for_message_prompt = "Wait for {other_agent} to send a message..."
+        self.wait_for_message_prompt = ""
+        self.last_message_prompt = "{other_agent} said: {last_message}"
+        # self.send_message_prompt = (
+        #     f"Send your message now (max {self.num_message_chars} chars)."
+        # )
+        self.send_message_prompt = f"Send your message now in <message>...</message> (<={self.num_message_chars} chars)."
+
+    def get_message_regex(self, observation: TrustAndSplitObs) -> str:
+        return rf"<message>[\s\S]{{0,{self.num_message_chars}}}</message>"
+
+    # def get_message_regex(self, observation: TrustAndSplitObs) -> str:
+    #     return rf"(?s).{{0,{self.num_message_chars}}}"
+
+    def get_split_regex(self, observation: TrustAndSplitObs) -> str:
+        items = list(observation.quantities.keys())
+        # Accept both singular and plural forms
+        item_pattern = "|".join(
+            [f"{item[:-1]}s?" if item.endswith("s") else f"{item}s?" for item in items]
+        )
+        regex = rf"(?i)<items_to_self> ?((?:\s*(?P<num>(10|[0-9]))\s*(?P<item>{item_pattern})\s*,?)+) ?</items_to_self>"
+        return regex
+
+    def get_split_action(
+        self, policy_output: str, observation: TrustAndSplitObs
+    ) -> Split:
+        items = list(observation.quantities.keys())
+        import re as _re
+
+        split_regex = self.get_split_regex(observation)
+        items_given_to_self = {item: 0 for item in items}
+        m = _re.match(split_regex, policy_output.strip())
+        if m:
+            # Find all (number, item) pairs
+            item_pattern = "|".join(
+                [
+                    f"{item[:-1]}s?" if item.endswith("s") else f"{item}s?"
+                    for item in items
+                ]
+            )
+            inner_regex = rf"(?i)(10|[0-9])\s*({item_pattern})"
+
+            def normalize_item_name(item_str):
+                for orig in items:
+                    if item_str.lower() == orig.lower():
+                        return orig
+                    if orig.endswith("s") and item_str.lower() == orig[:-1].lower():
+                        return orig
+                    if (
+                        not orig.endswith("s")
+                        and item_str.lower() == orig.lower() + "s"
+                    ):
+                        return orig
+
+            for num, item in _re.findall(inner_regex, m.group(1)):
+                items_given_to_self[normalize_item_name(item)] = int(num)
+        return Split(items_given_to_self=items_given_to_self)

src_code_for_reproducibility/markov_games/negotiation/tas_rps_simulation.py

@@ -0,0 +1,248 @@
+"""
+Trust-and-Split simulation.
+
+This environment models a simple bargaining game over 10 coins with messaging.
+Agents are assigned rock/paper/scissors hands, with the winner getting value 10 per coin
+and the loser getting value 1 per coin. Agents alternate sending messages for a fixed
+number of turns per round and then each submits a split proposal indicating how many
+coins they keep for themselves. Rewards are proportional if the proposed totals exceed 10.
+"""
+
+import copy
+from dataclasses import dataclass
+from typing import Any, Dict, List, Literal, Tuple
+
+from numpy.random import default_rng
+
+from mllm.markov_games.negotiation.nego_simulation import (
+    Message,
+    NegotiationObs,
+    NegotiationSimulation,
+    NegotiationState,
+    Split,
+    compute_tas_style_rewards,
+)
+from mllm.markov_games.rollout_tree import SimulationStepLog
+
+AgentId = str
+
+
+def _get_rps_winner(
+    hand1: Literal["rock", "paper", "scissors"],
+    hand2: Literal["rock", "paper", "scissors"],
+) -> Literal["rock", "paper", "scissors"]:
+    """Determine winner of rock-paper-scissors between two hands."""
+    if hand1 == hand2:
+        raise ValueError("Hands should be different")
+    if (
+        (hand1 == "rock" and hand2 == "scissors")
+        or (hand1 == "paper" and hand2 == "rock")
+        or (hand1 == "scissors" and hand2 == "paper")
+    ):
+        return hand1
+    else:
+        return hand2
+
+
+@dataclass
+class TrustAndSplitRPSState(NegotiationState):
+    hands: Dict[
+        AgentId, Literal["rock", "paper", "scissors"]
+    ]  # rock, paper, or scissors
+    previous_hands: Dict[AgentId, Literal["rock", "paper", "scissors"]] | None
+
+
+@dataclass
+class TrustAndSplitRPSObs(NegotiationObs):
+    hand: Literal["rock", "paper", "scissors"]
+    last_hand_agent: Literal["rock", "paper", "scissors"] | None
+    last_hand_coagent: Literal["rock", "paper", "scissors"] | None
+    last_hand_value_coagent: Literal["upper", "lower"] | None
+
+
+class TrustAndSplitRPSSimulation(NegotiationSimulation):
+    def __init__(
+        self,
+        alternating_hands: bool = False,
+        alternating_mix_ratio: float = None,
+        *args,
+        **kwargs,
+    ):
+        self.alternating_hands = alternating_hands
+        self.alternating_mix_ratio = alternating_mix_ratio
+        super().__init__(*args, **kwargs)
+        if self.alternating_mix_ratio is not None:
+            if self.rng.random() < self.alternating_mix_ratio:
+                self.alternating_hands = True
+            else:
+                self.alternating_hands = False
+
+    def _sample_hands_and_values(
+        self,
+        alternate_hands: bool = False,
+    ) -> Tuple[Dict[AgentId, str], Dict[AgentId, float]]:
+        hands = ["rock", "paper", "scissors"]
+        if alternate_hands:
+            previous_hands = list(self.state.previous_hands.values())
+            hand1, hand2 = self.rng.choice(hands, size=2, replace=False)
+            winner = _get_rps_winner(hand1, hand2)
+            loser = hand1 if winner == hand2 else hand2
+            previous_winner = _get_rps_winner(previous_hands[0], previous_hands[1])
+            agent_hands, values = {}, {}
+            for agent_id in self.agent_ids:
+                if self.state.previous_hands[agent_id] == previous_winner:
+                    agent_hands[agent_id] = loser
+                    values[agent_id] = 1.0
+                else:
+                    agent_hands[agent_id] = winner
+                    values[agent_id] = 10.0
+            return agent_hands, values
+        else:
+            # Assign different hands to each agent
+            hand1, hand2 = self.rng.choice(hands, size=2, replace=False)
+
+            agent_hands = {self.agent_ids[0]: hand1, self.agent_ids[1]: hand2}
+
+            # Determine winner and assign values
+            winner = _get_rps_winner(hand1, hand2)
+            values = {}
+            for agent_id in self.agent_ids:
+                if agent_hands[agent_id] == winner:
+                    values[agent_id] = 10.0  # Winner gets value 10
+                else:
+                    values[agent_id] = 1.0  # Loser gets value 1
+
+            return agent_hands, values
+
+    def set_new_round_of_variant(self):
+        self.state.previous_hands = copy.deepcopy(self.state.hands)
+        new_hands, new_values = self._sample_hands_and_values(
+            alternate_hands=self.alternating_hands
+        )
+        self.state.hands = new_hands
+        self.state.values = new_values
+        # Quantities are constant in TAS
+        self.state.quantities = {"coins": 10}
+        self.state.split_phase = False
+
+    def get_info_of_variant(
+        self, state: NegotiationState, actions: Dict[AgentId, Any]
+    ) -> Dict[str, Any]:
+        return {
+            "quantities": copy.deepcopy(state.quantities),
+            "hands": copy.deepcopy(state.hands),
+            "values": copy.deepcopy(state.values),
+            "previous_hands": copy.deepcopy(state.previous_hands),
+            "previous_values": copy.deepcopy(state.previous_values),
+            "splits": copy.deepcopy(state.splits),
+        }
+
+    def get_rewards(self, splits: Dict[AgentId, Split]) -> Dict[AgentId, float]:
+        return compute_tas_style_rewards(
+            self.agent_ids, self.state.values, splits, self.state.quantities
+        )
+
+    def get_obs_agent(self, agent_id):
+        """Returns observation for agent_id"""
+        other_id = self._other(agent_id)
+        last_value_coagent = (
+            None
+            if self.state.previous_values is None
+            else self.state.previous_values.get(other_id)
+        )
+        last_hand_coagent = (
+            None
+            if self.state.previous_hands is None
+            else self.state.previous_hands.get(other_id)
+        )
+        last_points_coagent = (
+            None
+            if self.state.previous_points is None
+            else round(self.state.previous_points.get(other_id), 1)
+        )
+        last_value_agent = (
+            None
+            if self.state.previous_values is None
+            else self.state.previous_values.get(agent_id)
+        )
+        last_hand_agent = (
+            None
+            if self.state.previous_hands is None
+            else self.state.previous_hands.get(agent_id)
+        )
+        last_points_agent = (
+            None
+            if self.state.previous_points is None
+            else round(self.state.previous_points.get(agent_id), 1)
+        )
+        last_split_coagent = None
+        last_split_agent = None
+        if self.state.previous_splits is not None:
+            last_split_coagent = self.state.previous_splits[
+                other_id
+            ].items_given_to_self["coins"]
+            last_split_agent = self.state.previous_splits[agent_id].items_given_to_self[
+                "coins"
+            ]
+        if last_hand_agent is None or last_hand_coagent is None:
+            last_hand_value_coagent = None
+        else:
+            winner = _get_rps_winner(last_hand_agent, last_hand_coagent)
+            last_hand_value_coagent = (
+                "upper" if winner == last_hand_coagent else "lower"
+            )
+        obs = TrustAndSplitRPSObs(
+            round_nb=self.state.round_nb,
+            last_message=self.state.last_message,
+            quota_messages_per_agent_per_round=self.quota_messages_per_agent_per_round,
+            current_agent=self.state.current_agent,
+            other_agent=self.agent_id_to_name[other_id],
+            quantities={"coins": 10},
+            item_types=self.item_types,
+            value=self.state.values[agent_id],
+            split_phase=self.state.split_phase,
+            last_split_agent=last_split_agent,
+            last_value_agent=last_value_agent,
+            last_points_agent=last_points_agent,
+            last_split_coagent=last_split_coagent,
+            last_value_coagent=last_value_coagent,
+            last_points_coagent=last_points_coagent,
+            hand=self.state.hands[agent_id],
+            last_hand_coagent=last_hand_coagent,
+            last_hand_agent=last_hand_agent,
+            last_quantities=self.state.previous_quantities,
+            last_hand_value_coagent=last_hand_value_coagent,
+        )
+        return obs
+
+    def get_state(self):
+        return self.state
+
+    def get_safe_copy(self):
+        """Return a safe copy of the simulation."""
+        simulation_copy = copy.copy(self)
+        simulation_copy.state = copy.deepcopy(self.state)
+        return simulation_copy
+
+    def reset(self):
+        """Initialize and return initial observations"""
+        # Decide starting agent alternating across resets for determinism
+        start_agent = self.agent_ids[self._starting_agent_index]
+        hands, values = self._sample_hands_and_values()
+        self.state = TrustAndSplitRPSState(
+            round_nb=0,
+            last_message="",
+            current_agent=start_agent,
+            quantities={"coins": 10},
+            values=values,
+            splits={aid: None for aid in self.agent_ids},
+            nb_messages_sent={aid: 0 for aid in self.agent_ids},
+            previous_values=None,
+            previous_splits=None,
+            previous_points=None,
+            split_phase=False,
+            hands=hands,
+            previous_hands=None,
+            previous_quantities=None,
+        )
+        return self.get_obs()

src_code_for_reproducibility/markov_games/negotiation/tas_simple_simulation.py

@@ -0,0 +1,169 @@
+import copy
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Dict, List, Literal
+
+from numpy.random import default_rng
+
+from mllm.markov_games.negotiation.nego_simulation import (
+    NegotiationObs,
+    NegotiationSimulation,
+    NegotiationState,
+    Split,
+    compute_tas_style_rewards,
+)
+
+AgentId = str
+
+
+@dataclass
+class TrustAndSplitSimpleState(NegotiationState):
+    pass
+
+
+@dataclass
+class TrustAndSplitSimpleObs(NegotiationObs):
+    last_value_str_coagent: str | None
+
+
+class TrustAndSplitSimpleSimulation(NegotiationSimulation):
+    def __init__(
+        self,
+        game_type: Literal["10-1-exclusive", "1-to-10"] = "1-to-10",
+        dist_type: Literal["uniform", "bimodal"] = "uniform",
+        beta_dist_alpha: float = 0.1,
+        beta_dist_beta: float = 0.1,
+        *args,
+        **kwargs,
+    ):
+        self.game_type = game_type
+        self.dist_type = dist_type
+        self.beta_dist_alpha = beta_dist_alpha
+        self.beta_dist_beta = beta_dist_beta
+        super().__init__(*args, **kwargs)
+
+    def _sample_values(self) -> Dict[AgentId, dict]:
+        values = {}
+        while True:
+            if self.game_type == "10-1-exclusive":
+                v = int(self.rng.choice([1, 10]))
+                values[self.agent_ids[0]] = v
+                values[self.agent_ids[1]] = 10 if v == 1 else 1
+            elif self.game_type == "1-to-10":
+                for aid in self.agent_ids:
+                    if self.dist_type == "uniform":
+                        values[aid] = int(self.rng.integers(1, 11))
+                    elif self.dist_type == "bimodal":
+                        alpha, beta = self.beta_dist_alpha, self.beta_dist_beta
+                        values[aid] = int(round(self.rng.beta(alpha, beta) * 9) + 1)
+            if len(set(values.values())) != 1:
+                break
+        return values
+
+    def _sample_quantities(self) -> Dict[str, int]:
+        return {"coins": 10}
+
+    def set_new_round_of_variant(self):
+        self.state.quantities = self._sample_quantities()
+        self.state.values = self._sample_values()
+        self.state.split_phase = False
+
+    def get_info_of_variant(
+        self, state: NegotiationState, actions: Dict[AgentId, Any]
+    ) -> Dict[str, Any]:
+        return {
+            "quantities": copy.deepcopy(state.quantities),
+            "values": copy.deepcopy(state.values),
+            # "previous_values": copy.deepcopy(state.previous_values),
+            "splits": copy.deepcopy(state.splits),
+        }
+
+    def get_rewards(self, splits: Dict[AgentId, Split]) -> Dict[AgentId, float]:
+        return compute_tas_style_rewards(
+            self.agent_ids, self.state.values, splits, self.state.quantities
+        )
+
+    def get_obs(self):
+        return {agent_id: self.get_obs_agent(agent_id) for agent_id in self.agent_ids}
+
+    def get_obs_agent(self, agent_id):
+        other_id = self._other(agent_id)
+        last_value_coagent = (
+            None
+            if self.state.previous_values is None
+            else self.state.previous_values.get(other_id)
+        )
+        last_points_coagent = (
+            None
+            if self.state.previous_points is None
+            else round(self.state.previous_points.get(other_id), 1)
+        )
+        last_value_agent = (
+            None
+            if self.state.previous_values is None
+            else self.state.previous_values.get(agent_id)
+        )
+        last_points_agent = (
+            None
+            if self.state.previous_points is None
+            else round(self.state.previous_points.get(agent_id), 1)
+        )
+        last_split_coagent = None
+        last_split_agent = None
+        if self.state.previous_splits is not None:
+            last_split_coagent = self.state.previous_splits[
+                other_id
+            ].items_given_to_self["coins"]
+            last_split_agent = self.state.previous_splits[agent_id].items_given_to_self[
+                "coins"
+            ]
+        if last_value_agent is None or last_value_coagent is None:
+            last_value_str_coagent = None
+        else:
+            if last_value_coagent > last_value_agent:
+                last_value_str_coagent = "higher"
+            elif last_value_coagent < last_value_agent:
+                last_value_str_coagent = "lower"
+            else:
+                raise ValueError("Should not be equal values")
+
+        obs = TrustAndSplitSimpleObs(
+            round_nb=self.state.round_nb,
+            last_message=self.state.last_message,
+            quota_messages_per_agent_per_round=self.quota_messages_per_agent_per_round,
+            current_agent=self.state.current_agent,
+            other_agent=self.agent_id_to_name[other_id],
+            quantities=self.state.quantities,
+            item_types=self.item_types,
+            value=self.state.values[agent_id],
+            split_phase=self.state.split_phase,
+            last_split_agent=last_split_agent,
+            last_value_agent=last_value_agent,
+            last_points_agent=last_points_agent,
+            last_split_coagent=last_split_coagent,
+            last_value_coagent=last_value_coagent,
+            last_points_coagent=last_points_coagent,
+            last_quantities=self.state.previous_quantities,
+            last_value_str_coagent=last_value_str_coagent,
+        )
+        return obs
+
+    def reset(self):
+        start_agent = self.agent_ids[self._starting_agent_index]
+        quantities = self._sample_quantities()
+        values = self._sample_values()
+        self.state = TrustAndSplitSimpleState(
+            round_nb=0,
+            last_message="",
+            current_agent=start_agent,
+            quantities=quantities,
+            values=values,
+            previous_values=None,
+            splits={aid: None for aid in self.agent_ids},
+            nb_messages_sent={aid: 0 for aid in self.agent_ids},
+            split_phase=False,
+            previous_splits=None,
+            previous_points=None,
+            previous_quantities=None,
+        )
+        return self.get_obs()

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

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

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", "<<message_start>>", "<<message_end>>"
+    elif "proposal_start" in regex and "proposal_end" in regex:
+        return "proposal", "<<proposal_start>>", "<<proposal_end>>"
+    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]

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:
+        ...

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.",
+            )

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"]

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]

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<think>\n([\s\S]*?)</think>\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,
+        )

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

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