Add files using upload-large-folder tool

src_code_for_reproducibility/chat_utils/apply_template.py

@@ -0,0 +1,84 @@
+import torch
+
+from mllm.chat_utils.chat_turn import ChatTurn
+from mllm.chat_utils.template_specific import (
+    custom_gemma3_template,
+    custom_llama3_template,
+    custom_qwen2_template,
+    custom_qwen3_template,
+    gemma3_assistant_postfix,
+    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
+    elif "gemma" in tokenizer.name_or_path.lower():
+        return custom_gemma3_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
+    elif "gemma" in tokenizer.name_or_path.lower():
+        return gemma3_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,109 @@
+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()
+)
+gemma3_assistant_postfix = (
+    AutoTokenizer.from_pretrained("google/gemma-3-4b-it")
+    .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 %}
+"""
+
+custom_gemma3_template = """
+{%- if add_system_prompt %}
+{{- bos_token -}}
+{%- endif %}
+{%- for message in messages -%}
+{%- if message['role'] == 'assistant' -%}
+{%- set role = 'model' -%}
+{%- else -%}
+{%- set role = message['role'] -%}
+{%- endif -%}
+{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+{{ '<start_of_turn>model\n' }}
+{%- 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/alternative_actions_runner.py

@@ -0,0 +1,138 @@
+import asyncio
+import copy
+import json
+import os.path
+from typing import Any, Tuple
+
+from mllm.markov_games.markov_game import AgentAndActionSafeCopy, MarkovGame
+from mllm.markov_games.rollout_tree import (
+    AgentActLog,
+    RolloutTreeBranchNode,
+    RolloutTreeNode,
+    RolloutTreeRootNode,
+    StepLog,
+)
+
+AgentId = str
+
+
+
+async def run_with_unilateral_alt_action(
+    markov_game: MarkovGame,
+    agent_id: AgentId,
+    time_step: int,
+    branch_node: RolloutTreeBranchNode,
+    max_depth: int,
+):
+    """
+    This function is used to generate a new branch for a given agent.
+    """
+
+    # Generate alternative action and take a step
+    await markov_game.set_action_of_agent(agent_id)
+    terminated: bool = markov_game.take_simulation_step()
+    step_log = markov_game.get_step_log()
+    first_alternative_node = RolloutTreeNode(
+        step_log=step_log,
+        time_step=time_step,
+    )
+
+    # Generate rest of trajectory up to max depth
+    time_step += 1
+    counter = 1
+    previous_node = first_alternative_node
+    while not terminated and counter <= max_depth:
+        terminated, step_log = await markov_game.step()
+        current_node = RolloutTreeNode(step_log=step_log, time_step=time_step)
+        previous_node.child = current_node
+        previous_node = current_node
+        counter += 1
+        time_step += 1
+
+    if branch_node.branches == None:
+        branch_node.branches = {agent_id: [first_alternative_node]}
+    else:
+        agent_branches = branch_node.branches.get(agent_id, [])
+        agent_branches.append(first_alternative_node)
+        branch_node.branches[agent_id] = agent_branches
+
+
+async def AlternativeActionsRunner(
+    markov_game: MarkovGame,
+    output_folder: str,
+    nb_alternative_actions: int,
+    max_depth: int,
+    branch_only_on_new_round: bool = False,
+):
+    """
+    This method generates a trajectory with partially completed branches,
+    where the branching comes from taking unilateraly different actions.
+    The resulting data is used to estimate the updated advantage alignment policy gradient terms.
+    Let k := nb_sub_steps. Then the number of steps generated is O(Tk), where T is
+    the maximum trajectory length.
+    """
+
+    tasks = []
+    time_step = 0
+    terminated = False
+    root = RolloutTreeRootNode(
+        id=markov_game.get_id(),
+        crn_id=markov_game.get_crn_id()
+    )
+    previous_node = root
+
+    while not terminated:
+        mg_before_action = markov_game.get_safe_copy()
+
+        # Get safe copies for main branch
+        agent_action_safe_copies: dict[
+            AgentId, AgentAndActionSafeCopy
+        ] = await markov_game.get_actions_of_agents_without_side_effects()
+
+        markov_game.set_actions_of_agents_manually(agent_action_safe_copies)
+        terminated = markov_game.take_simulation_step()
+        main_node = RolloutTreeNode(
+            step_log=markov_game.get_step_log(), time_step=time_step
+        )
+        branch_node = RolloutTreeBranchNode(main_child=main_node)
+        previous_node.child = branch_node
+        previous_node = main_node
+
+        # Get alternative branches by generating new unilateral actions
+        for agent_id in markov_game.agent_ids:
+            for _ in range(nb_alternative_actions):
+                # Get safe copies for branches
+                branch_agent_action_safe_copies: dict[
+                    AgentId, AgentAndActionSafeCopy
+                ] = {
+                    agent_id: AgentAndActionSafeCopy(
+                        action=copy.deepcopy(agent_action_safe_copy.action),
+                        action_info=copy.deepcopy(agent_action_safe_copy.action_info),
+                        agent_after_action=agent_action_safe_copy.agent_after_action.get_safe_copy(),
+                    )
+                    for agent_id, agent_action_safe_copy in agent_action_safe_copies.items()
+                }
+                mg_branch: MarkovGame = mg_before_action.get_safe_copy()
+                other_agent_id = [id for id in mg_branch.agent_ids if id != agent_id][0]
+                mg_branch.set_action_and_agent_after_action_manually(
+                    agent_id=other_agent_id,
+                    agent_action_safe_copy=branch_agent_action_safe_copies[
+                        other_agent_id
+                    ],
+                )
+                task = asyncio.create_task(
+                    run_with_unilateral_alt_action(
+                        markov_game=mg_branch,
+                        time_step=time_step,
+                        agent_id=agent_id,
+                        branch_node=branch_node,
+                        max_depth=max_depth,
+                    )
+                )
+                tasks.append(task)
+        time_step += 1
+
+    # wait for all branches to complete
+    await asyncio.gather(*tasks)
+
+    return root

src_code_for_reproducibility/markov_games/group_timesteps.py

@@ -0,0 +1,150 @@
+"""
+This module contains the logic for grouping time steps.
+"""
+import copy
+from typing import Callable
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import (
+    AgentActLog,
+    RolloutTreeBranchNode,
+    RolloutTreeNode,
+    RolloutTreeRootNode,
+    StepLog,
+)
+from mllm.markov_games.simulation import SimulationStepLog
+
+AgentId = str
+
+
+def group_time_steps(
+    rollout_tree: RolloutTreeRootNode,
+    accumulation_stop_condition: Callable[[StepLog], bool],
+) -> RolloutTreeRootNode:
+    """
+    During generation, we create rollout trees according to the real time steps.
+    However, during training, we might want to treat groups of time steps as a single time step.
+    As a concrete example, take Trust-and-Split. At each round, say we have X time steps of communication and then one time step for the split.
+    Then the communication actions will not get any reward, and the split action will get the reward. During REINFORCE training, with discounting, this
+    can cause training instability. We could instead treat every action in the round as being part of a single action, and give it the reward of the split action.
+    This method helps to do this sort of grouping.
+    It accumulates actions until the accumulation_stop_condition is met, and then creates a new node with the accumulated actions.
+    It then recursively calls itself on the child node.
+    Details:
+    - The reward for the group is the reward of the last time step in the group.
+    - The simulation log for the group is the simulation log of the last time step in the group.
+    - The state end for the group becomes the first state end in the group.
+    - The agent info for the group is the agent info of the last time step in the group.
+    """
+
+    def group_step_logs(step_logs: list[StepLog]) -> StepLog:
+        """
+        Concatenate per-agent chat turns across steps; keep only the first is_state_end.
+        """
+        last_sim_log = step_logs[-1].simulation_step_log
+        agent_ids = {aid for s in step_logs for aid in s.action_logs.keys()}
+        grouped_logs: dict[AgentId, AgentActLog] = {}
+        for aid in agent_ids:
+            turns = []
+            for s in step_logs:
+                act = s.action_logs.get(aid)
+                if act and act.chat_turns:
+                    turns.extend(copy.deepcopy(act.chat_turns))
+            disable_is_state_end = False
+            # Only the first state_end should be True, the rest should be False
+            for t in turns:
+                if t.is_state_end:
+                    if disable_is_state_end:
+                        t.is_state_end = False
+                    else:
+                        disable_is_state_end = True
+                    continue
+            grouped_logs[aid] = AgentActLog(
+                chat_turns=turns, info=step_logs[-1].action_logs[aid].info
+            )
+        return StepLog(action_logs=grouped_logs, simulation_step_log=last_sim_log)
+
+    def group_time_steps_rec(
+        current_node: RolloutTreeNode | RolloutTreeBranchNode,
+        group_time_step: int,
+        accumulation_step_logs: list[StepLog],
+    ) -> RolloutTreeNode | RolloutTreeBranchNode:
+        """
+        Groups time steps. Recursion is used to handle branches.
+        """
+        assert isinstance(current_node, RolloutTreeNode) or isinstance(
+            current_node, RolloutTreeBranchNode
+        ), "Current node must be a tree node or a branch node. Is of type: " + str(
+            type(current_node)
+        )
+        first_group_node = None
+        current_group_node = None
+        while current_node is not None:
+            if isinstance(current_node, RolloutTreeBranchNode):
+                raise Exception(
+                    "Grouping timesteps by round is not supported for branching trajectories yet."
+                )
+            # Special recursive case for branches
+            # if isinstance(current_node, RolloutTreeBranchNode):
+            #     branches = {}
+            #     for agent_id, branch_nodes in current_node.branches.items():
+            #         branch_group_nodes = []
+            #         for branch_node in branch_nodes:
+            #             branch_group_node = group_time_steps_rec(
+            #                 current_node=branch_node,
+            #                 group_time_step=group_time_step,
+            #                 accumulation_step_logs=copy.deepcopy(accumulation_step_logs))
+            #             branch_group_nodes.append(branch_group_node)
+            #         branches[agent_id] = branch_group_nodes
+
+            #     main_child_group_node = group_time_steps_rec(
+            #         current_node=current_node.main_child,
+            #         group_time_step=group_time_step,
+            #         accumulation_step_logs=copy.deepcopy(accumulation_step_logs))
+
+            #     return RolloutTreeBranchNode(main_child=main_child_group_node, branches=branches)
+
+            # Accumulate
+            accumulation_step_logs.append(current_node.step_log)
+            if accumulation_stop_condition(current_node.step_log):
+                grouped_step_logs = group_step_logs(accumulation_step_logs)
+                accumulation_step_logs = []
+                new_group_node = RolloutTreeNode(
+                    step_log=grouped_step_logs, time_step=group_time_step, child=None
+                )
+                if first_group_node == None:
+                    first_group_node = new_group_node
+                group_time_step += 1
+                if current_group_node is not None:
+                    current_group_node.child = new_group_node
+                current_group_node = new_group_node
+            current_node = current_node.child
+        return first_group_node
+
+    node = group_time_steps_rec(
+        current_node=rollout_tree.child, group_time_step=0, accumulation_step_logs=[]
+    )
+    return RolloutTreeRootNode(
+        id=rollout_tree.id,
+        crn_id=rollout_tree.crn_id,
+        child=node,
+        agent_ids=rollout_tree.agent_ids,
+    )
+
+
+def stop_when_round_ends(step_log: StepLog) -> bool:
+    """
+    Simplest stop condition. Will return True if step log is the last time step of a round.
+    This will throw an error if this information is not available in the simulation info.
+    """
+    assert (
+        "is_last_timestep_in_round" in step_log.simulation_step_log.info.keys()
+    ), "To group by round, is_last_timestep_in_round must be set in the info of your simulation step log at each time step."
+    return step_log.simulation_step_log.info["is_last_timestep_in_round"]
+
+
+def group_by_round(rollout_tree: RolloutTreeRootNode) -> RolloutTreeRootNode:
+    """
+    Groups time steps by round.
+    """
+    return group_time_steps(rollout_tree, stop_when_round_ends)

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/markov_game.py

@@ -0,0 +1,208 @@
+"""
+This class unifies a simulation, and the agents acting in it (see `simulation.py` & `agent.py`).
+In a MarkovGame step,
+    1) each agent takes an action,
+    2) the state transitions with respect to these actions,
+    3) all relevant data of the step is appended to the historical data list
+
+In order to perform 3), the agents and the simulation are expected, at each time step,
+to return a log of the state transition (from their perspective).
+For instance, the Simulation might send rewards and the agents might send prompting contexts to be used later to generate the training data.
+A different approach would be to simply have the agents keep their data private and log it upon completion of a trajectory.
+The approach we use here centralizes the data gathering aspect,
+making it easy to create sub-trajectories (in the `runners` defined in `runners.py`) descriptions that
+only log information for step transitions occuring after the branching out.
+"""
+import asyncio
+import copy
+import json
+import os
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, Tuple
+
+from transformers.models.idefics2 import Idefics2Config
+
+from mllm.markov_games.agent import Agent
+from mllm.markov_games.rollout_tree import AgentActLog, StepLog
+from mllm.markov_games.simulation import Simulation
+
+AgentId = str
+
+
+@dataclass
+class AgentAndActionSafeCopy:
+    action: Any
+    action_info: AgentActLog
+    agent_after_action: type[Agent]
+
+
+class MarkovGame(object):
+    def __init__(
+        self,
+        id: int,
+        agents: dict[AgentId, type[Agent]],
+        simulation: type[Simulation],
+        crn_id: int,
+    ):
+        """
+        Args:
+            agents:
+            output_path:
+                Path where the step infos are saved.
+            simulation:
+                Simulation object. Example: IPDSimulation
+        """
+        self.agents = agents
+        self.agent_ids = self.agents.keys()
+        self.simulation = simulation
+        self.simulation_step_log = None
+        self.agent_step_logs = {agent_id: None for agent_id in self.agent_ids}
+        self.actions = {}
+        self.id = id
+        self.crn_id = crn_id
+
+    def get_id(self) -> str:
+        return self.id
+
+    def get_crn_id(self) -> int:
+        return self.crn_id
+
+    def get_agent_ids(self) -> List[AgentId]:
+        return list(self.agent_ids)
+
+    async def get_action_of_agent_without_side_effects(
+        self, agent_id: AgentId
+    ) -> Tuple[Any, AgentActLog]:
+        """
+        Safe function to get an action of an agent without modifying the agent or the simulation.
+        """
+        agent = self.agents[agent_id]
+        agent_before_action = agent.get_safe_copy()
+        obs = self.simulation.get_obs_agent(agent_id)
+        action, action_info = await agent.act(observation=obs)
+        self.agents[agent_id] = agent_before_action
+        agent_after_action = agent.get_safe_copy()
+        return AgentAndActionSafeCopy(action, action_info, agent_after_action)
+
+    async def get_actions_of_agents_without_side_effects(
+        self,
+    ) -> dict[AgentId, AgentAndActionSafeCopy]:
+        """
+        Safe function to get an action of an agent without modifying the agent or the simulation.
+        """
+        tasks = []
+        for agent_id in self.agent_ids:
+            task = asyncio.create_task(
+                self.get_action_of_agent_without_side_effects(agent_id)
+            )
+            tasks.append(task)
+        agent_and_action_safe_copies: list[
+            AgentAndActionSafeCopy
+        ] = await asyncio.gather(*tasks)
+        return {
+            agent_id: agent_and_action_safe_copy
+            for agent_id, agent_and_action_safe_copy in zip(
+                self.agent_ids, agent_and_action_safe_copies
+            )
+        }
+
+    def set_action_and_agent_after_action_manually(
+        self,
+        agent_id: AgentId,
+        agent_action_safe_copy: AgentAndActionSafeCopy,
+    ):
+        """
+        Set the action and the agent after action manually.
+        """
+        self.actions[agent_id] = agent_action_safe_copy.action
+        self.agent_step_logs[agent_id] = agent_action_safe_copy.action_info
+        self.agents[agent_id] = agent_action_safe_copy.agent_after_action
+
+    def set_actions_of_agents_manually(
+        self, actions: dict[AgentId, AgentAndActionSafeCopy]
+    ):
+        """
+        Set the actions of agents manually.
+        """
+        for agent_id, agent_action_safe_copy in actions.items():
+            self.set_action_and_agent_after_action_manually(
+                agent_id, agent_action_safe_copy
+            )
+
+    async def set_action_of_agent(self, agent_id: AgentId):
+        """
+        TOWRITE
+        """
+        agent = self.agents[agent_id]
+        obs = self.simulation.get_obs_agent(agent_id)
+        action, action_info = await agent.act(observation=obs)
+        self.actions[agent_id] = action
+        self.agent_step_logs[agent_id] = action_info
+
+    async def set_actions(self):
+        """
+        TOWRITE
+        """
+        # background_tasks = set()
+        tasks = []
+        for agent_id in self.agent_ids:
+            task = asyncio.create_task(self.set_action_of_agent(agent_id))
+            tasks.append(task)
+        await asyncio.gather(*tasks)
+
+    def take_simulation_step(self):
+        """
+        TOWRITE
+        """
+        terminated, self.simulation_step_log = self.simulation.step(self.actions)
+        return terminated
+
+    def get_step_log(self) -> StepLog:
+        """
+        TOWRITE
+        TODO: assert actions and simulation have taken step
+        """
+        step_log = StepLog(
+            simulation_step_log=self.simulation_step_log,
+            action_logs=self.agent_step_logs,
+        )
+        return step_log
+
+    async def step(self) -> Tuple[bool, StepLog]:
+        """
+        TOWRITE
+        """
+        await self.set_actions()
+        terminated = self.take_simulation_step()
+        step_log = self.get_step_log()
+        return terminated, step_log
+
+    def get_safe_copy(self):
+        """
+        TOWRITE
+        """
+
+        new_markov_game = copy.copy(self)
+        new_simulation = self.simulation.get_safe_copy()
+        new_agents = {
+            agent_id: agent.get_safe_copy() for agent_id, agent in self.agents.items()
+        }
+
+        # Reassign copied components
+        new_markov_game.simulation = new_simulation
+        new_markov_game.agents = new_agents
+
+        # IMPORTANT: ensure agent_ids references the new agents dict, not the original
+        new_markov_game.agent_ids = new_markov_game.agents.keys()
+
+        # Deep-copy step data to avoid correlation
+        new_markov_game.simulation_step_log = copy.deepcopy(self.simulation_step_log)
+        new_markov_game.actions = copy.deepcopy(self.actions)
+        # Rebuild logs to align exactly with new agent ids
+        old_agent_step_logs = copy.deepcopy(self.agent_step_logs)
+        new_markov_game.agent_step_logs = {
+            agent_id: old_agent_step_logs.get(agent_id)
+            for agent_id in new_markov_game.agent_ids
+        }
+
+        return new_markov_game

src_code_for_reproducibility/markov_games/mg_utils.py

@@ -0,0 +1,89 @@
+import asyncio
+import copy
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from mllm.markov_games.ipd.ipd_agent import IPDAgent
+from mllm.markov_games.ipd.ipd_simulation import IPD
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.negotiation.dond_agent import DealNoDealAgent
+from mllm.markov_games.negotiation.dond_simulation import DealNoDealSimulation
+from mllm.markov_games.negotiation.nego_hard_coded_policies import (
+    HardCodedNegoGreedyPolicy,
+    HardCodedNegoWelfareMaximizingPolicy,
+)
+from mllm.markov_games.ipd.Ipd_hard_coded_agents import AlwaysCooperateIPDAgent, AlwaysDefectIPDAgent
+from mllm.markov_games.negotiation.no_press_nego_agent import NoPressAgent
+from mllm.markov_games.negotiation.no_press_nego_simulation import NoPressSimulation
+from mllm.markov_games.negotiation.tas_agent import TrustAndSplitAgent
+from mllm.markov_games.negotiation.tas_rps_agent import TrustAndSplitRPSAgent
+from mllm.markov_games.negotiation.tas_rps_simulation import TrustAndSplitRPSSimulation
+from mllm.markov_games.negotiation.tas_simple_agent import TrustAndSplitSimpleAgent
+from mllm.markov_games.negotiation.tas_simple_simulation import (
+    TrustAndSplitSimpleSimulation,
+)
+from mllm.markov_games.negotiation.tas_simulation import TrustAndSplitSimulation
+from mllm.markov_games.rollout_tree import (
+    AgentActLog,
+    RolloutTreeBranchNode,
+    RolloutTreeNode,
+    RolloutTreeRootNode,
+    StepLog,
+)
+from mllm.markov_games.simulation import SimulationStepLog
+
+AgentId = str
+
+
+@dataclass
+class AgentConfig:
+    agent_id: str
+    agent_name: str
+    agent_class_name: str
+    policy_id: str
+    init_kwargs: dict
+
+
+@dataclass
+class MarkovGameConfig:
+    id: int
+    seed: int
+    simulation_class_name: str
+    simulation_init_args: dict
+    agent_configs: list[AgentConfig]
+
+
+def init_markov_game_components(
+    config: MarkovGameConfig, policies: dict[str, Callable[[list[dict]], str]]
+):
+    """
+    TOWRITE
+    """
+    agents = {}
+    agent_names = []
+    for agent_config in config.agent_configs:
+        agent_id = agent_config.agent_id
+        agent_name = agent_config.agent_name
+        agent_class = eval(agent_config.agent_class_name)
+        agent = agent_class(
+            seed=config.seed,
+            agent_id=agent_id,
+            agent_name=agent_name,
+            policy=policies[agent_config.policy_id],
+            **agent_config.init_kwargs,
+        )
+        agents[agent_id] = agent
+        agent_names.append(agent_name)
+    simulation = eval(config.simulation_class_name)(
+        seed=config.seed,
+        agent_ids=list(agents.keys()),
+        agent_names=agent_names,
+        **config.simulation_init_args,
+    )
+    markov_game = MarkovGame(
+        id=config.id,
+        crn_id=config.seed,
+        agents=agents,
+        simulation=simulation,
+    )
+    return markov_game

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_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/negotiation_statistics.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+from typing import Callable, Dict, List, Tuple
+
+from mllm.markov_games.negotiation.nego_simulation import Split
+from mllm.markov_games.rollout_tree import SimulationStepLog
+
+
+def avg_reward(sl: SimulationStepLog) -> List[Tuple[str, float]]:
+    """Average (per-step) reward for each agent and overall.
+
+    What it computes:
+            - Returns the raw reward for every (non-buffer) agent at the current
+                simulation step.
+            - Adds an aggregate key ``all_agents`` which is the simple arithmetic
+                mean across the agents present in ``sl.rewards``.
+
+    Rationale / motivation:
+            Monitoring the reward stream at each step helps:
+                * Diagnose reward shaping issues (e.g., unintended negative drift).
+                * Provide a fairness snapshot (are rewards systematically skewed?).
+                * Supply a ubiquitous baseline metric used by other higher‑level
+                    summaries (efficiency, surplus allocation, etc.).
+
+    Return shape:
+            { agent_id: float, ..., "all_agents": float }
+            If any agent id contains the substring "buffer" we treat this step as
+            an implementation artifact (e.g., rollout buffer) and return ``None``
+            to avoid polluting aggregates.
+    """
+    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]
+
+
+def split_efficiency(sl: SimulationStepLog) -> List[Tuple[str, float]] | None:
+    """Final‑round allocation efficiency relative to an upper bound.
+
+    What it computes (only on the last timestep of a negotiation round):
+            - Uses ``info['values']`` (per‑agent per‑item valuations) and
+                ``info['quantities']`` (available item counts) to form a greedy
+                *upper bound* on achievable total reward: allocate each unit of an
+                item to the single agent who values that item most.
+            - Compares the actually realized sum of rewards at that final
+                timestep to this constructed maximum.
+            - Emits a single scalar under key ``"all_agents"`` equal to
+                achieved / theoretical_max.
+
+    Motivation:
+            Efficiency (a core welfare notion) distinguishes between coordination
+            failures (low efficiency) versus strategic distributional disputes
+            (high efficiency but uneven splits). Tracking this per round helps
+            evaluate whether models learn to identify and realize joint surplus.
+
+    Notes / caveats:
+            - Only defined for 2+ non‑buffer agents; if a buffer agent is present
+                returns ``None`` to exclude spurious steps.
+            - Requires the environment to have populated ``values`` and
+                ``quantities``; otherwise returns ``None``.
+            - This is an optimistic bound (not necessarily reachable under
+                protocol constraints) but is simple, fast, and comparable across
+                runs.
+    """
+    info = sl.info or {}
+    if not info or not info.get("is_last_timestep_in_round"):
+        return None
+    quantities = info.get("quantities") or {}
+    values = info.get("values") or {}
+    if not values or not quantities:
+        return None
+    agent_ids = list(sl.rewards.keys())
+    if type(values[agent_ids[0]]) is dict:
+        item_keys = list(values.values())[0].keys()
+        max_vals, max_quantities = [], []
+        for item in item_keys:
+            max_val = max(float(agent_vals[item]) for agent_vals in values.values())
+            max_vals.append(max_val)
+            max_quantities.append(quantities[item])
+    else:
+        max_vals = [max(float(v) for v in values.values())]
+        max_quantities = [quantities[item] for item in quantities.keys()]
+    for aid in sl.rewards.keys():
+        if "buffer" in str(aid) and "live" not in str(aid):
+            return None
+    achieved = sum(float(v) for v in sl.rewards.values())
+    max_reward = sum(d * v for d, v in zip(max_quantities, max_vals))
+    # Efficiency is a global metric; emit same value for a special key "all"
+    return [("split_efficiency", achieved / max_reward)]
+
+
+def _extract_items_from_split(raw_split: Dict) -> Dict[str, float] | None:
+    """Return a mapping item->proposal amount from a split structure.
+
+    Supports both generic negotiation splits with nested structure
+    { 'items_given_to_self': {item: qty, ...}}
+    and TAS coin-only variants which may already be a flat mapping {'coins': qty}.
+    """
+
+    if raw_split is None:
+        return {}
+    elif isinstance(raw_split, Split):
+        return {k: float(v) for k, v in raw_split.items_given_to_self.items()}
+    elif isinstance(raw_split, dict):
+        if "items_given_to_self" in raw_split and isinstance(
+            raw_split["items_given_to_self"], dict
+        ):
+            return {k: float(v) for k, v in raw_split["items_given_to_self"].items()}
+        # Fallback: assume already flat mapping of items
+        elif hasattr(raw_split, "items_given_to_self"):
+            return {k: float(v) for k, v in raw_split["items_given_to_self"].items()}
+        return {
+            k: float(v) for k, v in raw_split.items() if isinstance(v, (int, float))
+        }
+    return {}
+
+
+def _average_proposal_relative_value(
+    sl: SimulationStepLog,
+    metric_name: str,
+    comparator: Callable[[float, float], bool],
+    opposite_comparator: Callable[[float, float], bool],
+) -> Dict[str, float | None] | None:
+    """Shared implementation for proposal size conditioned on relative value.
+
+    Parameters:
+            comparator: returns True when agent_0's value relation (e.g. < or >)
+                                    to agent_1 holds for an item and we should collect agent_0's
+                                    proposed quantity for that item.
+            opposite_comparator: inverse relation used to collect agent_1's items.
+
+    Behavior:
+            - Executes only on final timestep of a round (where the definitive
+                proposal / allocation is known via ``info['splits']``).
+            - For each item, classifies which agent's value satisfies the chosen
+                relation and records that agent's proposed quantity from the split.
+            - Averages (mean) across all qualifying items per agent; if no items
+                qualify for an agent returns ``None`` for that agent id.
+            - Adds ``all_agents`` mean across the numeric (non-None) agent values.
+
+    Why this matters:
+            Distinguishing how much an agent *asks for* when it subjectively
+            values items more (or less) than its counterpart reveals patterns of
+            opportunism vs. concession. This is especially useful when raw reward
+            differences are subtle but allocation *intent* differs.
+    """
+    info = sl.info or {}
+    if not info or not info.get("is_last_timestep_in_round"):
+        return None
+    quantities = info.get("quantities") or {}
+    splits = info.get("splits") or {}
+    values = info.get("values") or {}
+    agent_ids: List[str] = list(sl.rewards.keys())
+    if len(agent_ids) != 2:
+        return None  # Only defined for 2-agent case.
+    for aid in agent_ids:
+        if "buffer" in str(aid) and "live" not in str(aid):
+            return None
+    # Extract per-agent item proposals robustly
+    split_items = {aid: _extract_items_from_split(splits.get(aid)) for aid in agent_ids}
+    agent_0_vals: List[float] = []
+    agent_1_vals: List[float] = []
+    for item in quantities.keys():
+        # Values may be either a float (same for all items) or dict per item
+        v0_raw = values[agent_ids[0]]
+        v1_raw = values[agent_ids[1]]
+        v0 = float(v0_raw[item]) if isinstance(v0_raw, dict) else float(v0_raw)
+        v1 = float(v1_raw[item]) if isinstance(v1_raw, dict) else float(v1_raw)
+        if comparator(v0, v1):
+            agent_0_vals.append(split_items[agent_ids[0]].get(item, 0.0))
+        elif opposite_comparator(v0, v1):
+            agent_1_vals.append(split_items[agent_ids[1]].get(item, 0.0))
+    out: Dict[str, float | None] = {}
+    out[f"{metric_name}-{agent_ids[0]}"] = (
+        sum(agent_0_vals) / len(agent_0_vals) if agent_0_vals else None
+    )
+    out[f"{metric_name}-{agent_ids[1]}"] = (
+        sum(agent_1_vals) / len(agent_1_vals) if agent_1_vals else None
+    )
+
+    return [(key, value) for key, value in out.items() if value is not None]
+
+
+def average_proposal_when_agent_values_item_lower(
+    sl: SimulationStepLog,
+) -> List[Tuple[str, float | None]] | None:
+    """Mean quantity an agent proposes for items it values *less* than opponent.
+
+    Interpretation:
+        A higher value implies the agent still claims (or is allocated) a
+        notable share of items where it has a comparative *disadvantage* in
+        valuation, signaling either strategic over-claiming or protocol-driven
+        egalitarian splits. Conversely, very low numbers can indicate
+        efficient specialization or excessive concession.
+
+    Returns:
+        Mapping { agent_id: float | None, "all_agents": float | None } where
+        None indicates no qualifying items for that agent in the round.
+    """
+    return _average_proposal_relative_value(
+        sl,
+        "average_proposal_when_agent_values_item_lower",
+        lambda a, b: a < b,
+        lambda a, b: a > b,
+    )
+
+
+def average_proposal_when_agent_values_item_higher(
+    sl: SimulationStepLog,
+) -> List[Tuple[str, float | None]] | None:
+    """Mean quantity an agent proposes for items it values *more* than opponent.
+
+    Interpretation:
+        Captures how aggressively an agent claims items where it holds a
+        comparative *advantage*. Elevated values can reflect rational
+        specialization (efficient exploitation of comparative advantage) or
+        potentially unfair grabs if paired with low concession in the lower
+        valuation metric. Comparing this with the 'lower' counterpart helps
+        profile negotiation style (cooperative vs. exploitative).
+
+    Returns:
+        Mapping { agent_id: float | None, "all_agents": float | None } where
+        None indicates no qualifying items.
+    """
+    return _average_proposal_relative_value(
+        sl,
+        "average_proposal_when_agent_values_item_higher",
+        lambda a, b: a > b,
+        lambda a, b: a < b,
+    )
+
+
+# Explicit list of metric functions exported for rendering. Helper functions
+# starting with '_' are intentionally excluded. Update this list when adding
+# new public statistics so render.py can rely on it instead of introspecting
+# every callable in the module.
+stat_functs: list[Callable[[SimulationStepLog], List[Tuple[str, float]]]] = [
+    avg_reward,
+    average_proposal_when_agent_values_item_lower,
+    average_proposal_when_agent_values_item_higher,
+    split_efficiency,
+]

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_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_simulation.py

@@ -0,0 +1,172 @@
+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 TrustAndSplitState(NegotiationState):
+    pass
+
+
+@dataclass
+class TrustAndSplitObs(NegotiationObs):
+    pass
+
+
+class TrustAndSplitSimulation(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))
+            agent_values = [sum(v.values()) for v in values.values()]
+            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
+                    ):  # Different values for this item
+                        has_conflict = True
+                        break
+                if not has_conflict:
+                    continue
+            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 = 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
+            last_split_agent = self.state.previous_splits[agent_id].items_given_to_self
+        obs = TrustAndSplitObs(
+            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,
+        )
+        return obs
+
+    def reset(self):
+        start_agent = self.agent_ids[self._starting_agent_index]
+        quantities = self._sample_quantities()
+        values = self._sample_values()
+        self.state = TrustAndSplitState(
+            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/rollout_tree.py

@@ -0,0 +1,86 @@
+"""
+TODO: add parent to nodes so that some verification can be done. For instance, to ensure that node reward keys match the parent node.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, List, Literal, Optional, Tuple
+
+import jsonschema
+from pydantic import BaseModel, Field, model_validator
+
+from mllm.chat_utils.chat_turn import ChatTurn
+
+AgentId = str
+
+
+class SimulationStepLog(BaseModel):
+    rewards: dict[AgentId, float]
+    info: Any = None
+
+
+class AgentActLog(BaseModel):
+    chat_turns: list[ChatTurn] | None
+    info: Any = None
+
+    @model_validator(mode="after")
+    def _exactly_one_state_end(self):
+        """
+        This method is used to enforce that for each AgentActLog, there is exactly one ChatTurn which is a state end.
+        """
+        if self.chat_turns != []:
+            n = sum(1 for t in self.chat_turns if t.is_state_end)
+            if n != 1:
+                raise ValueError(
+                    f"AgentActLog must have exactly one ChatTurn with is_state_end=True; got {self.chat_turns}."
+                )
+            return self
+        else:
+            return self
+
+
+class StepLog(BaseModel):
+    action_logs: dict[AgentId, AgentActLog]
+    simulation_step_log: SimulationStepLog
+
+
+# BranchType = Literal["unilateral_deviation", "common_deviation"] # might not be necessary
+# class BranchNodeInfo(BaseModel):
+#     branch_id: str
+#     branch_for: AgentId
+#     branch_type: BranchType
+
+
+class RolloutTreeNode(BaseModel):
+    step_log: StepLog
+    time_step: int
+    child: RolloutTreeNode | RolloutTreeBranchNode | None = None
+
+
+class RolloutTreeBranchNode(BaseModel):
+    """
+    First item of the tuple indicates which agent "called" for an alternative branch.
+    """
+
+    main_child: RolloutTreeNode
+    branches: dict[AgentId, list[RolloutTreeNode]] | None = None
+
+
+class RolloutTreeRootNode(BaseModel):
+    id: int
+    crn_id: int  # ID of the rng used to generate this rollout tree
+    child: RolloutTreeNode | RolloutTreeBranchNode | None = None
+    agent_ids: List[AgentId] = Field(min_length=1)
+
+
+# class RolloutTreeLeafNode(BaseModel):
+#     step_log: StepLog
+#     time_step: int
+
+
+# Necessary for self-referential stuff in pydantic
+RolloutTreeBranchNode.model_rebuild()
+RolloutTreeNode.model_rebuild()

src_code_for_reproducibility/markov_games/run_markov_games.py

@@ -0,0 +1,24 @@
+import asyncio
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from torch._C import ClassType
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import RolloutTreeRootNode
+
+
+async def run_markov_games(
+    runner: Callable[[MarkovGame], RolloutTreeRootNode],
+    runner_kwargs: dict,
+    output_folder: str,
+    markov_games: list[MarkovGame],
+) -> list[RolloutTreeRootNode]:
+    tasks = []
+    for mg in markov_games:
+        tasks.append(
+            asyncio.create_task(
+                runner(markov_game=mg, output_folder=output_folder, **runner_kwargs)
+            )
+        )
+    return await asyncio.gather(*tasks)

src_code_for_reproducibility/markov_games/simulation.py

@@ -0,0 +1,87 @@
+"""
+A Simulation is the environment of a Markov Game.
+The Simulation is not responsible for properly checking / formatting the responses of LLM's.
+This is the job of the `Agent` class.
+Simulations expect clean actions, and are defined similarly to `gymnasium` environments, except that they are adapted for the Multi-agent setting.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Tuple
+
+from numpy.random import default_rng
+
+from mllm.markov_games.rollout_tree import SimulationStepLog
+
+
+class Simulation(ABC):
+    @abstractmethod
+    def __init__(self, seed: int, *args, **kwargs):
+        self.seed = seed
+        self.rng = default_rng(self.seed)
+
+    @abstractmethod
+    def step(self, actions: Any) -> Tuple[bool, SimulationStepLog]:
+        """
+        Returns terminated, info
+        """
+        raise NotImplementedError
+
+    def get_obs(self):
+        """Returns all agent observations in dict
+
+        Returns:
+            observations
+        """
+        raise NotImplementedError
+
+    def get_obs_agent(self, agent_id):
+        """Returns observation for agent_id"""
+        raise NotImplementedError
+
+    def get_obs_size(self):
+        """Returns the shape of the observation"""
+        raise NotImplementedError
+
+    def get_state(self):
+        raise NotImplementedError
+
+    def get_state_size(self):
+        """Returns the shape of the state"""
+        raise NotImplementedError
+
+    def get_avail_actions(self):
+        raise NotImplementedError
+
+    def get_avail_agent_actions(self, agent_id):
+        """Returns the available actions for agent_id"""
+        raise NotImplementedError
+
+    def get_total_actions(self):
+        """Returns the total number of actions an agent could ever take"""
+        # TODO: This is only suitable for a discrete 1 dimensional action space for each agent
+        raise NotImplementedError
+
+    def get_safe_copy(self):
+        """
+        Return copy of the agent object that is decorrelated from the original object.
+        """
+        raise NotImplementedError
+
+    def reset(self):
+        """Returns initial observations and states"""
+        raise NotImplementedError
+
+    def render(self):
+        raise NotImplementedError
+
+    def close(self):
+        raise NotImplementedError
+
+    # def seed(self):
+    #     raise NotImplementedError
+
+    def save_replay(self):
+        raise NotImplementedError
+
+    def get_simulation_info(self):
+        raise NotImplementedError

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