diff --git a/src_code_for_reproducibility/__pycache__/__init__.cpython-311.pyc b/src_code_for_reproducibility/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2b817647d28cfd5c3fd8ea2da23692ec7b1aaacc
Binary files /dev/null and b/src_code_for_reproducibility/__pycache__/__init__.cpython-311.pyc differ
diff --git a/src_code_for_reproducibility/__pycache__/__init__.cpython-312.pyc b/src_code_for_reproducibility/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e1bb76ab4b9c5bb44943e5eabe2ca3dadde518d2
Binary files /dev/null and b/src_code_for_reproducibility/__pycache__/__init__.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/chat_utils/apply_template.py b/src_code_for_reproducibility/chat_utils/apply_template.py
new file mode 100644
index 0000000000000000000000000000000000000000..f143ab15e6b30637700df3381b7de686153017cd
--- /dev/null
+++ b/src_code_for_reproducibility/chat_utils/apply_template.py
@@ -0,0 +1,78 @@
+import torch
+
+from mllm.chat_utils.chat_turn import ChatTurn
+from mllm.chat_utils.template_specific import (
+ custom_llama3_template,
+ custom_qwen2_template,
+ custom_qwen3_template,
+ qwen2_assistant_postfix,
+ qwen3_assistant_postfix,
+)
+
+
+def get_custom_chat_template(tokenizer) -> str:
+ """
+ Get the chat template for the tokenizer.
+ """
+ if "qwen2" in tokenizer.name_or_path.lower():
+ return custom_qwen2_template
+ elif "llama" in tokenizer.name_or_path.lower():
+ return custom_llama3_template
+ elif "qwen3" in tokenizer.name_or_path.lower():
+ return custom_qwen3_template
+ else:
+ raise ValueError(f"Tokenizer {tokenizer.name_or_path} not supported")
+
+
+def get_custom_assistant_postfix(tokenizer) -> torch.Tensor:
+ """
+ Get the custom assistant postfix for the tokenizer.
+ """
+ if "qwen2" in tokenizer.name_or_path.lower():
+ return qwen2_assistant_postfix
+ elif "qwen3" in tokenizer.name_or_path.lower():
+ return qwen3_assistant_postfix
+ return torch.tensor([], dtype=torch.long)
+
+
+def tokenize_chats(chats: list[ChatTurn], tokenizer, enable_thinking) -> None:
+ """
+ Set the chat_template_token_ids for each chat turn.
+ # TODO: use engine tokens if available
+ """
+ custom_template = get_custom_chat_template(tokenizer)
+ custom_assistant_postfix: torch.Tensor = get_custom_assistant_postfix(tokenizer)
+ for i, chat in enumerate(chats):
+ if chat.chat_template_token_ids is None:
+ if chat.role == "user":
+ next_chat = chats[i + 1] if i + 1 < len(chats) else None
+ add_generation_prompt = True
+ if next_chat and next_chat.role == "user":
+ add_generation_prompt = False
+ encoded_chat = tokenizer.apply_chat_template(
+ [chat],
+ return_tensors="pt",
+ chat_template=custom_template,
+ add_generation_prompt=add_generation_prompt,
+ add_system_prompt=True if i == 0 else False,
+ enable_thinking=enable_thinking,
+ ).flatten()
+ previous_chat = chats[i - 1] if i > 0 else None
+ if previous_chat and previous_chat.role == "assistant":
+ encoded_chat = torch.cat([custom_assistant_postfix, encoded_chat])
+ elif chat.role == "assistant":
+ encoded_chat = chat.out_token_ids
+ chat.chat_template_token_ids = encoded_chat
+
+
+def chat_turns_to_token_ids(
+ chats: list[ChatTurn], tokenizer, enable_thinking
+) -> list[int]:
+ """
+ Tokenize the chat turns and set the chat_template_token_ids for each chat turn.
+ """
+ tokenize_chats(chats=chats, tokenizer=tokenizer, enable_thinking=enable_thinking)
+ token_ids = []
+ for chat in chats:
+ token_ids.append(chat.chat_template_token_ids)
+ return torch.cat(token_ids)
diff --git a/src_code_for_reproducibility/chat_utils/chat_turn.py b/src_code_for_reproducibility/chat_utils/chat_turn.py
new file mode 100644
index 0000000000000000000000000000000000000000..d22c1b853def374fa0ee21eb523f2b9d104ed35b
--- /dev/null
+++ b/src_code_for_reproducibility/chat_utils/chat_turn.py
@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, List, Literal, Optional, Tuple
+
+import jsonschema
+import torch
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+AgentId = str
+
+
+class ChatTurn(BaseModel):
+ model_config = ConfigDict(arbitrary_types_allowed=True) # needed for torch tensors
+
+ role: str = Field(pattern="^(user|assistant)$")
+ agent_id: AgentId # ID of the agent with which the chat occured
+ content: str
+ reasoning_content: str | None = None
+ chat_template_token_ids: torch.LongTensor | None = None # Token ids of chat template format. For example, token ids of "{content}""
+ out_token_ids: torch.LongTensor | None = (
+ None # tokens generated from inference engine
+ )
+ log_probs: torch.FloatTensor | None = None
+ is_state_end: bool = False # indicates whether this chat turn marks the end of a state in the trajectory
diff --git a/src_code_for_reproducibility/chat_utils/template_specific.py b/src_code_for_reproducibility/chat_utils/template_specific.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d8f8e5f1be6bcafd1b9e1a465c446253b985e3e
--- /dev/null
+++ b/src_code_for_reproducibility/chat_utils/template_specific.py
@@ -0,0 +1,87 @@
+import huggingface_hub
+import torch
+from transformers import AutoTokenizer
+
+custom_llama3_template = """
+{%- if add_system_prompt %}
+ {{- '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nCutting Knowledge Date: December 2023\nToday Date: 26 Jul 2024\n\n<|eot_id|>' }}
+{%- endif %}
+{%- for message in messages %}
+ {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+ {{- '<|start_header_id|>' + 'assistant' + '<|end_header_id|>\n\n' }}
+{%- endif %}
+"""
+
+qwen2_assistant_postfix = (
+ AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
+ .encode("\n", return_tensors="pt")
+ .flatten()
+)
+qwen3_assistant_postfix = (
+ AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
+ .encode("\n", return_tensors="pt")
+ .flatten()
+)
+custom_qwen2_template = """
+{%- if add_system_prompt %}
+ {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages %}
+ {%- if message.content is string %}
+ {%- set content = message.content %}
+ {%- else %}
+ {%- set content = '' %}
+ {%- endif %}
+ {%- if (message.role == "user") %}
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+ {%- elif message.role == "assistant" %}
+ {%- set reasoning_content = '' %}
+ {%- if message.reasoning_content is string %}
+ {%- set reasoning_content = message.reasoning_content %}
+ {%- else %}
+ {%- if '' in content %}
+ {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %}
+ {%- set content = content.split('')[-1].lstrip('\n') %}
+ {%- endif %}
+ {%- endif %}
+ {%- if loop.index0 > ns.last_query_index %}
+ {%- if reasoning_content %}
+ {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }}
+ {%- else %}
+ {{- '<|im_start|>' + message.role + '\n' + content }}
+ {%- endif %}
+ {%- else %}
+ {{- '<|im_start|>' + message.role + '\n' + content }}
+ {%- endif %}
+ {{- '<|im_end|>\n' }}
+ {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+ {{- '<|im_start|>assistant\n' }}
+{%- endif %}
+"""
+
+custom_qwen3_template = """
+{%- for message in messages %}
+ {%- if message.content is string %}
+ {%- set content = message.content %}
+ {%- else %}
+ {%- set content = '' %}
+ {%- endif %}
+ {%- if (message.role == "user") %}
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+ {%- elif message.role == "assistant" %}
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+ {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+ {{- '<|im_start|>assistant\n' }}
+ {%- if enable_thinking is defined and enable_thinking is false %}
+ {{- '\n\n\n\n' }}
+ {%- endif %}
+{%- endif %}
+"""
diff --git a/src_code_for_reproducibility/docs/Makefile b/src_code_for_reproducibility/docs/Makefile
new file mode 100644
index 0000000000000000000000000000000000000000..9375655022388557d65910465b09d2a5fa9d2e4a
--- /dev/null
+++ b/src_code_for_reproducibility/docs/Makefile
@@ -0,0 +1,19 @@
+# Minimal makefile for Sphinx documentation
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS ?=
+SPHINXBUILD ?= sphinx-build
+SOURCEDIR = source
+BUILDDIR = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+ @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+ @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS)
diff --git a/src_code_for_reproducibility/docs/generate_docs.py b/src_code_for_reproducibility/docs/generate_docs.py
new file mode 100644
index 0000000000000000000000000000000000000000..e644cbbf091a97420500fb47346c07be5ed141ac
--- /dev/null
+++ b/src_code_for_reproducibility/docs/generate_docs.py
@@ -0,0 +1,249 @@
+#!/usr/bin/env python3
+"""
+Script to automatically generate Sphinx documentation for all modules and build the HTML website.
+"""
+import importlib.util
+import os
+import subprocess
+import sys
+
+
+def check_and_install_dependencies():
+ """Check for required dependencies and install them if missing."""
+ required_packages = [
+ "sphinx",
+ "sphinx-rtd-theme",
+ "sphinxcontrib-napoleon",
+ "sphinxcontrib-mermaid",
+ "sphinx-autodoc-typehints",
+ ]
+
+ missing_packages = []
+
+ for package in required_packages:
+ # Convert package name to module name (replace - with _)
+ module_name = package.replace("-", "_")
+
+ # Check if the package is installed
+ if importlib.util.find_spec(module_name) is None:
+ missing_packages.append(package)
+
+ # Install missing packages
+ if missing_packages:
+ print(f"Installing missing dependencies: {', '.join(missing_packages)}")
+ subprocess.check_call(
+ [sys.executable, "-m", "pip", "install"] + missing_packages
+ )
+ print("Dependencies installed successfully")
+ else:
+ print("All required dependencies are already installed")
+
+
+def create_makefile(docs_dir):
+ """Create a Makefile for Sphinx documentation if it doesn't exist."""
+ makefile_path = os.path.join(docs_dir, "Makefile")
+
+ if os.path.exists(makefile_path):
+ print(f"Makefile already exists at {makefile_path}")
+ return
+
+ print(f"Creating Makefile at {makefile_path}")
+
+ makefile_content = """# Minimal makefile for Sphinx documentation
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS ?=
+SPHINXBUILD ?= sphinx-build
+SOURCEDIR = source
+BUILDDIR = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+ @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+ @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS)
+"""
+
+ with open(makefile_path, "w") as f:
+ f.write(makefile_content)
+
+ print("Makefile created successfully")
+
+
+def create_make_bat(docs_dir):
+ """Create a make.bat file for Windows if it doesn't exist."""
+ make_bat_path = os.path.join(docs_dir, "make.bat")
+
+ if os.path.exists(make_bat_path):
+ print(f"make.bat already exists at {make_bat_path}")
+ return
+
+ print(f"Creating make.bat at {make_bat_path}")
+
+ make_bat_content = """@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+ set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+ echo.
+ echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+ echo.installed, then set the SPHINXBUILD environment variable to point
+ echo.to the full path of the 'sphinx-build' executable. Alternatively you
+ echo.may add the Sphinx directory to PATH.
+ echo.
+ echo.If you don't have Sphinx installed, grab it from
+ echo.https://www.sphinx-doc.org/
+ exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd
+"""
+
+ with open(make_bat_path, "w") as f:
+ f.write(make_bat_content)
+
+ print("make.bat created successfully")
+
+
+def main():
+ # Check and install required dependencies
+ print("=== Checking dependencies ===")
+ check_and_install_dependencies()
+
+ # Get the directory of this script
+ script_dir = os.path.dirname(os.path.abspath(__file__))
+
+ # Path to the project root
+ project_root = os.path.dirname(script_dir)
+
+ # Path to the source directory
+ source_dir = os.path.join(project_root, "src")
+
+ # Path to the docs source directory
+ docs_source_dir = os.path.join(script_dir, "source")
+
+ # Print paths for debugging
+ print(f"Script directory: {script_dir}")
+ print(f"Project root: {project_root}")
+ print(f"Source directory: {source_dir}")
+ print(f"Docs source directory: {docs_source_dir}")
+
+ # Make sure the source directory exists
+ if not os.path.exists(source_dir):
+ print(f"Error: Source directory {source_dir} does not exist!")
+ sys.exit(1)
+
+ # Make sure the docs source directory exists
+ if not os.path.exists(docs_source_dir):
+ print(f"Creating docs source directory: {docs_source_dir}")
+ os.makedirs(docs_source_dir)
+
+ # Step 1: Run sphinx-apidoc to generate .rst files for all modules
+ print("\n=== Generating API documentation ===")
+ cmd = [
+ "sphinx-apidoc",
+ "-f", # Force overwriting of existing files
+ "-e", # Put module documentation before submodule documentation
+ "-M", # Put module documentation before subpackage documentation
+ "-o",
+ docs_source_dir, # Output directory
+ source_dir, # Source code directory
+ ]
+
+ print(f"Running command: {' '.join(cmd)}")
+ result = subprocess.run(cmd, capture_output=True, text=True)
+
+ # Print the output of the command
+ print("STDOUT:")
+ print(result.stdout)
+
+ print("STDERR:")
+ print(result.stderr)
+
+ if result.returncode != 0:
+ print(f"Error: sphinx-apidoc failed with return code {result.returncode}")
+ sys.exit(1)
+
+ # List the files in the docs source directory
+ print("\nFiles in docs/source directory:")
+ for file in sorted(os.listdir(docs_source_dir)):
+ print(f" {file}")
+
+ print("\nDocumentation source files generated successfully!")
+
+ # Step 2: Create Makefile and make.bat if they don't exist
+ create_makefile(script_dir)
+ create_make_bat(script_dir)
+
+ # Step 3: Build the HTML documentation
+ print("\n=== Building HTML documentation ===")
+
+ # Determine the build command based on the platform
+ if os.name == "nt": # Windows
+ build_cmd = ["make.bat", "html"]
+ else: # Unix/Linux/Mac
+ build_cmd = ["make", "html"]
+
+ # Change to the docs directory to run the build command
+ os.chdir(script_dir)
+
+ print(f"Running command: {' '.join(build_cmd)}")
+ build_result = subprocess.run(build_cmd, capture_output=True, text=True)
+
+ # Print the output of the build command
+ print("STDOUT:")
+ print(build_result.stdout)
+
+ print("STDERR:")
+ print(build_result.stderr)
+
+ if build_result.returncode != 0:
+ print(f"Error: HTML build failed with return code {build_result.returncode}")
+ sys.exit(1)
+
+ # Get the path to the built HTML documentation
+ html_dir = os.path.join(script_dir, "build", "html")
+ index_path = os.path.join(html_dir, "index.html")
+
+ if os.path.exists(index_path):
+ print(f"\nHTML documentation built successfully!")
+ print(f"You can view it by opening: {index_path}")
+
+ # Try to open the documentation in a browser
+ try:
+ import webbrowser
+
+ print("\nAttempting to open documentation in your default browser...")
+ webbrowser.open(f"file://{index_path}")
+ except Exception as e:
+ print(f"Could not open browser automatically: {e}")
+ else:
+ print(f"\nWarning: HTML index file not found at {index_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src_code_for_reproducibility/docs/make.bat b/src_code_for_reproducibility/docs/make.bat
new file mode 100644
index 0000000000000000000000000000000000000000..dc1312ab09ca6fb0267dee6b28a38e69c253631a
--- /dev/null
+++ b/src_code_for_reproducibility/docs/make.bat
@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+ set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+ echo.
+ echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+ echo.installed, then set the SPHINXBUILD environment variable to point
+ echo.to the full path of the 'sphinx-build' executable. Alternatively you
+ echo.may add the Sphinx directory to PATH.
+ echo.
+ echo.If you don't have Sphinx installed, grab it from
+ echo.https://www.sphinx-doc.org/
+ exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd
diff --git a/src_code_for_reproducibility/markov_games/__init__.py b/src_code_for_reproducibility/markov_games/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/__init__.cpython-311.pyc b/src_code_for_reproducibility/markov_games/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..87d439427c824e863dc4f7216198d6170539d30e
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/__init__.cpython-311.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/__init__.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5a3a6c09c6f7608acfeddebc65fdb4149e8c3cb7
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/__init__.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/agent.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/agent.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..efe246c29687cefb7069674e57114ea38dab252a
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/agent.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/alternative_actions_runner.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/alternative_actions_runner.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5788779130aac0fc02a9c1b2ac985b4cd5427eb1
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/alternative_actions_runner.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/gather_and_export_utils.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/gather_and_export_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ebad9ae7dd0fb04490740075b6d765f0d0237b2c
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/gather_and_export_utils.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/group_timesteps.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/group_timesteps.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..68ef22b6f0c70af174960ab268da175e89c9d7e3
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/group_timesteps.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/linear_runner.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/linear_runner.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1d4da97bca809a1e8526f5dc00c3ee23d8b5c021
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/linear_runner.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/markov_game.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/markov_game.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c9be34b4a5f90d0e3c622100495120366881e0b0
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/markov_game.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/mg_utils.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/mg_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2fc3c1e2c56e0ba4c7ea7d628eee3b00ad878f6a
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/mg_utils.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/rollout_tree.cpython-311.pyc b/src_code_for_reproducibility/markov_games/__pycache__/rollout_tree.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..56bfac1adc30d8ee8d1e0defc19d26ead0758f4b
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/rollout_tree.cpython-311.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/rollout_tree.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/rollout_tree.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bede63b1b249138a212198e9d178e47b63db35d2
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/rollout_tree.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/run_markov_games.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/run_markov_games.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..caebf44fc211bb439d6ffe745a8fe8b08d71f568
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/run_markov_games.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/__pycache__/simulation.cpython-312.pyc b/src_code_for_reproducibility/markov_games/__pycache__/simulation.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1b3d76acd3764faebf2875304319edea4aa5e856
Binary files /dev/null and b/src_code_for_reproducibility/markov_games/__pycache__/simulation.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/markov_games/agent.py b/src_code_for_reproducibility/markov_games/agent.py
new file mode 100644
index 0000000000000000000000000000000000000000..9961d860fecca7eaa4d13cb195b50dc658be5bec
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/agent.py
@@ -0,0 +1,76 @@
+"""
+In simple RL paradise, where the action dimensions are constant and well defined,
+Agent classes are not necessary. But in MARL, with LLM's, there isn't always
+a direct path from policy to action. For instance, from the observation of the environment,
+a prompt must be created. Then, the outputs of the policy might be incorrect, so a second
+request to the LLM must be sent before the action is well defined. This is why this Agent class exists.
+It acts as a mini environment, bridging the gap between the core simulation and
+the LLM policies.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Tuple
+
+from numpy.random import default_rng
+
+from mllm.markov_games.rollout_tree import AgentActLog
+
+
+class Agent(ABC):
+ @abstractmethod
+ def __init__(
+ self,
+ seed: int,
+ agent_id: str,
+ agent_name: str,
+ agent_policy: Callable[[list[dict]], str],
+ *args,
+ **kwargs,
+ ):
+ """
+ Initialize the agent state.
+ """
+ self.seed = seed
+ self.agent_id = agent_id
+ self.agent_name = agent_name
+ self.policy = policy
+ self.rng = default_rng(self.seed)
+ raise NotImplementedError
+
+ async def act(self, observation) -> Tuple[Any, AgentActLog]:
+ """
+ Query (possibly multiple times) a policy (or possibly a pool of policies) to
+ obtain the action of the agent.
+
+ Example:
+ action = None
+ prompt = self.observation_to_prompt(observation)
+ while not self.valid(action):
+ output = await self.policy.generate(prompt)
+ action = self.policy_output_to_action(output)
+ return action
+
+ Returns:
+ action
+ step_info
+ """
+ raise NotImplementedError
+
+ def get_safe_copy(self):
+ """
+ Return copy of the agent object that is decorrelated from the original object.
+ """
+ raise NotImplementedError
+
+ def reset(self):
+ raise NotImplementedError
+
+ def render(self):
+ raise NotImplementedError
+
+ def close(self):
+ raise NotImplementedError
+
+ def get_agent_info(self):
+ raise NotImplementedError
diff --git a/src_code_for_reproducibility/markov_games/alternative_actions_runner.py b/src_code_for_reproducibility/markov_games/alternative_actions_runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..c64db2deda539a1a71e045309cfdf257d2cbc614
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/alternative_actions_runner.py
@@ -0,0 +1,138 @@
+import asyncio
+import copy
+import json
+import os.path
+from typing import Any, Tuple
+
+from mllm.markov_games.markov_game import AgentAndActionSafeCopy, MarkovGame
+from mllm.markov_games.rollout_tree import (
+ AgentActLog,
+ RolloutTreeBranchNode,
+ RolloutTreeNode,
+ RolloutTreeRootNode,
+ StepLog,
+)
+
+AgentId = str
+
+
+
+async def run_with_unilateral_alt_action(
+ markov_game: MarkovGame,
+ agent_id: AgentId,
+ time_step: int,
+ branch_node: RolloutTreeBranchNode,
+ max_depth: int,
+):
+ """
+ This function is used to generate a new branch for a given agent.
+ """
+
+ # Generate alternative action and take a step
+ await markov_game.set_action_of_agent(agent_id)
+ terminated: bool = markov_game.take_simulation_step()
+ step_log = markov_game.get_step_log()
+ first_alternative_node = RolloutTreeNode(
+ step_log=step_log,
+ time_step=time_step,
+ )
+
+ # Generate rest of trajectory up to max depth
+ time_step += 1
+ counter = 1
+ previous_node = first_alternative_node
+ while not terminated and counter <= max_depth:
+ terminated, step_log = await markov_game.step()
+ current_node = RolloutTreeNode(step_log=step_log, time_step=time_step)
+ previous_node.child = current_node
+ previous_node = current_node
+ counter += 1
+ time_step += 1
+
+ if branch_node.branches == None:
+ branch_node.branches = {agent_id: [first_alternative_node]}
+ else:
+ agent_branches = branch_node.branches.get(agent_id, [])
+ agent_branches.append(first_alternative_node)
+ branch_node.branches[agent_id] = agent_branches
+
+
+async def AlternativeActionsRunner(
+ markov_game: MarkovGame,
+ output_folder: str,
+ nb_alternative_actions: int,
+ max_depth: int,
+ branch_only_on_new_round: bool = False,
+):
+ """
+ This method generates a trajectory with partially completed branches,
+ where the branching comes from taking unilateraly different actions.
+ The resulting data is used to estimate the updated advantage alignment policy gradient terms.
+ Let k := nb_sub_steps. Then the number of steps generated is O(Tk), where T is
+ the maximum trajectory length.
+ """
+
+ tasks = []
+ time_step = 0
+ terminated = False
+ root = RolloutTreeRootNode(
+ id=markov_game.get_id(),
+ crn_id=markov_game.get_crn_id()
+ )
+ previous_node = root
+
+ while not terminated:
+ mg_before_action = markov_game.get_safe_copy()
+
+ # Get safe copies for main branch
+ agent_action_safe_copies: dict[
+ AgentId, AgentAndActionSafeCopy
+ ] = await markov_game.get_actions_of_agents_without_side_effects()
+
+ markov_game.set_actions_of_agents_manually(agent_action_safe_copies)
+ terminated = markov_game.take_simulation_step()
+ main_node = RolloutTreeNode(
+ step_log=markov_game.get_step_log(), time_step=time_step
+ )
+ branch_node = RolloutTreeBranchNode(main_child=main_node)
+ previous_node.child = branch_node
+ previous_node = main_node
+
+ # Get alternative branches by generating new unilateral actions
+ for agent_id in markov_game.agent_ids:
+ for _ in range(nb_alternative_actions):
+ # Get safe copies for branches
+ branch_agent_action_safe_copies: dict[
+ AgentId, AgentAndActionSafeCopy
+ ] = {
+ agent_id: AgentAndActionSafeCopy(
+ action=copy.deepcopy(agent_action_safe_copy.action),
+ action_info=copy.deepcopy(agent_action_safe_copy.action_info),
+ agent_after_action=agent_action_safe_copy.agent_after_action.get_safe_copy(),
+ )
+ for agent_id, agent_action_safe_copy in agent_action_safe_copies.items()
+ }
+ mg_branch: MarkovGame = mg_before_action.get_safe_copy()
+ other_agent_id = [id for id in mg_branch.agent_ids if id != agent_id][0]
+ mg_branch.set_action_and_agent_after_action_manually(
+ agent_id=other_agent_id,
+ agent_action_safe_copy=branch_agent_action_safe_copies[
+ other_agent_id
+ ],
+ )
+ task = asyncio.create_task(
+ run_with_unilateral_alt_action(
+ markov_game=mg_branch,
+ time_step=time_step,
+ agent_id=agent_id,
+ branch_node=branch_node,
+ max_depth=max_depth,
+ )
+ )
+ tasks.append(task)
+ time_step += 1
+
+ # wait for all branches to complete
+ await asyncio.gather(*tasks)
+
+ return root
diff --git a/src_code_for_reproducibility/markov_games/group_timesteps.py b/src_code_for_reproducibility/markov_games/group_timesteps.py
new file mode 100644
index 0000000000000000000000000000000000000000..dad5271c500f539f2719110bd676f183746e51e4
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/group_timesteps.py
@@ -0,0 +1,150 @@
+"""
+This module contains the logic for grouping time steps.
+"""
+import copy
+from typing import Callable
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import (
+ AgentActLog,
+ RolloutTreeBranchNode,
+ RolloutTreeNode,
+ RolloutTreeRootNode,
+ StepLog,
+)
+from mllm.markov_games.simulation import SimulationStepLog
+
+AgentId = str
+
+
+def group_time_steps(
+ rollout_tree: RolloutTreeRootNode,
+ accumulation_stop_condition: Callable[[StepLog], bool],
+) -> RolloutTreeRootNode:
+ """
+ During generation, we create rollout trees according to the real time steps.
+ However, during training, we might want to treat groups of time steps as a single time step.
+ As a concrete example, take Trust-and-Split. At each round, say we have X time steps of communication and then one time step for the split.
+ Then the communication actions will not get any reward, and the split action will get the reward. During REINFORCE training, with discounting, this
+ can cause training instability. We could instead treat every action in the round as being part of a single action, and give it the reward of the split action.
+ This method helps to do this sort of grouping.
+ It accumulates actions until the accumulation_stop_condition is met, and then creates a new node with the accumulated actions.
+ It then recursively calls itself on the child node.
+ Details:
+ - The reward for the group is the reward of the last time step in the group.
+ - The simulation log for the group is the simulation log of the last time step in the group.
+ - The state end for the group becomes the first state end in the group.
+ - The agent info for the group is the agent info of the last time step in the group.
+ """
+
+ def group_step_logs(step_logs: list[StepLog]) -> StepLog:
+ """
+ Concatenate per-agent chat turns across steps; keep only the first is_state_end.
+ """
+ last_sim_log = step_logs[-1].simulation_step_log
+ agent_ids = {aid for s in step_logs for aid in s.action_logs.keys()}
+ grouped_logs: dict[AgentId, AgentActLog] = {}
+ for aid in agent_ids:
+ turns = []
+ for s in step_logs:
+ act = s.action_logs.get(aid)
+ if act and act.chat_turns:
+ turns.extend(copy.deepcopy(act.chat_turns))
+ disable_is_state_end = False
+ # Only the first state_end should be True, the rest should be False
+ for t in turns:
+ if t.is_state_end:
+ if disable_is_state_end:
+ t.is_state_end = False
+ else:
+ disable_is_state_end = True
+ continue
+ grouped_logs[aid] = AgentActLog(
+ chat_turns=turns, info=step_logs[-1].action_logs[aid].info
+ )
+ return StepLog(action_logs=grouped_logs, simulation_step_log=last_sim_log)
+
+ def group_time_steps_rec(
+ current_node: RolloutTreeNode | RolloutTreeBranchNode,
+ group_time_step: int,
+ accumulation_step_logs: list[StepLog],
+ ) -> RolloutTreeNode | RolloutTreeBranchNode:
+ """
+ Groups time steps. Recursion is used to handle branches.
+ """
+ assert isinstance(current_node, RolloutTreeNode) or isinstance(
+ current_node, RolloutTreeBranchNode
+ ), "Current node must be a tree node or a branch node. Is of type: " + str(
+ type(current_node)
+ )
+ first_group_node = None
+ current_group_node = None
+ while current_node is not None:
+ if isinstance(current_node, RolloutTreeBranchNode):
+ raise Exception(
+ "Grouping timesteps by round is not supported for branching trajectories yet."
+ )
+ # Special recursive case for branches
+ # if isinstance(current_node, RolloutTreeBranchNode):
+ # branches = {}
+ # for agent_id, branch_nodes in current_node.branches.items():
+ # branch_group_nodes = []
+ # for branch_node in branch_nodes:
+ # branch_group_node = group_time_steps_rec(
+ # current_node=branch_node,
+ # group_time_step=group_time_step,
+ # accumulation_step_logs=copy.deepcopy(accumulation_step_logs))
+ # branch_group_nodes.append(branch_group_node)
+ # branches[agent_id] = branch_group_nodes
+
+ # main_child_group_node = group_time_steps_rec(
+ # current_node=current_node.main_child,
+ # group_time_step=group_time_step,
+ # accumulation_step_logs=copy.deepcopy(accumulation_step_logs))
+
+ # return RolloutTreeBranchNode(main_child=main_child_group_node, branches=branches)
+
+ # Accumulate
+ accumulation_step_logs.append(current_node.step_log)
+ if accumulation_stop_condition(current_node.step_log):
+ grouped_step_logs = group_step_logs(accumulation_step_logs)
+ accumulation_step_logs = []
+ new_group_node = RolloutTreeNode(
+ step_log=grouped_step_logs, time_step=group_time_step, child=None
+ )
+ if first_group_node == None:
+ first_group_node = new_group_node
+ group_time_step += 1
+ if current_group_node is not None:
+ current_group_node.child = new_group_node
+ current_group_node = new_group_node
+ current_node = current_node.child
+ return first_group_node
+
+ node = group_time_steps_rec(
+ current_node=rollout_tree.child, group_time_step=0, accumulation_step_logs=[]
+ )
+ return RolloutTreeRootNode(
+ id=rollout_tree.id,
+ crn_id=rollout_tree.crn_id,
+ child=node,
+ agent_ids=rollout_tree.agent_ids,
+ )
+
+
+def stop_when_round_ends(step_log: StepLog) -> bool:
+ """
+ Simplest stop condition. Will return True if step log is the last time step of a round.
+ This will throw an error if this information is not available in the simulation info.
+ """
+ assert (
+ "is_last_timestep_in_round" in step_log.simulation_step_log.info.keys()
+ ), "To group by round, is_last_timestep_in_round must be set in the info of your simulation step log at each time step."
+ return step_log.simulation_step_log.info["is_last_timestep_in_round"]
+
+
+def group_by_round(rollout_tree: RolloutTreeRootNode) -> RolloutTreeRootNode:
+ """
+ Groups time steps by round.
+ """
+ return group_time_steps(rollout_tree, stop_when_round_ends)
diff --git a/src_code_for_reproducibility/markov_games/linear_runner.py b/src_code_for_reproducibility/markov_games/linear_runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..f81ab6d37ae9c680b2f2a53388988117d37f8a47
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/linear_runner.py
@@ -0,0 +1,30 @@
+import asyncio
+import json
+import os.path
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import RolloutTreeNode, RolloutTreeRootNode
+
+
+async def LinearRunner(
+ markov_game: MarkovGame, output_folder: str
+) -> RolloutTreeRootNode:
+ """
+ This method generates a trajectory without branching.
+ """
+ time_step = 0
+ terminated = False
+ root = RolloutTreeRootNode(
+ id=markov_game.get_id(),
+ crn_id=markov_game.get_crn_id(),
+ agent_ids=markov_game.get_agent_ids(),
+ )
+ previous_node = root
+ while not terminated:
+ terminated, step_log = await markov_game.step()
+ current_node = RolloutTreeNode(step_log=step_log, time_step=time_step)
+ previous_node.child = current_node
+ previous_node = current_node
+ time_step += 1
+
+ return root
diff --git a/src_code_for_reproducibility/markov_games/markov_game.py b/src_code_for_reproducibility/markov_games/markov_game.py
new file mode 100644
index 0000000000000000000000000000000000000000..73a48213bddcf0a59976fa0870eec19f59ae47d9
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/markov_game.py
@@ -0,0 +1,208 @@
+"""
+This class unifies a simulation, and the agents acting in it (see `simulation.py` & `agent.py`).
+In a MarkovGame step,
+ 1) each agent takes an action,
+ 2) the state transitions with respect to these actions,
+ 3) all relevant data of the step is appended to the historical data list
+
+In order to perform 3), the agents and the simulation are expected, at each time step,
+to return a log of the state transition (from their perspective).
+For instance, the Simulation might send rewards and the agents might send prompting contexts to be used later to generate the training data.
+A different approach would be to simply have the agents keep their data private and log it upon completion of a trajectory.
+The approach we use here centralizes the data gathering aspect,
+making it easy to create sub-trajectories (in the `runners` defined in `runners.py`) descriptions that
+only log information for step transitions occuring after the branching out.
+"""
+import asyncio
+import copy
+import json
+import os
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, Tuple
+
+from transformers.models.idefics2 import Idefics2Config
+
+from mllm.markov_games.agent import Agent
+from mllm.markov_games.rollout_tree import AgentActLog, StepLog
+from mllm.markov_games.simulation import Simulation
+
+AgentId = str
+
+
+@dataclass
+class AgentAndActionSafeCopy:
+ action: Any
+ action_info: AgentActLog
+ agent_after_action: type[Agent]
+
+
+class MarkovGame(object):
+ def __init__(
+ self,
+ id: int,
+ agents: dict[AgentId, type[Agent]],
+ simulation: type[Simulation],
+ crn_id: int,
+ ):
+ """
+ Args:
+ agents:
+ output_path:
+ Path where the step infos are saved.
+ simulation:
+ Simulation object. Example: IPDSimulation
+ """
+ self.agents = agents
+ self.agent_ids = self.agents.keys()
+ self.simulation = simulation
+ self.simulation_step_log = None
+ self.agent_step_logs = {agent_id: None for agent_id in self.agent_ids}
+ self.actions = {}
+ self.id = id
+ self.crn_id = crn_id
+
+ def get_id(self) -> str:
+ return self.id
+
+ def get_crn_id(self) -> int:
+ return self.crn_id
+
+ def get_agent_ids(self) -> List[AgentId]:
+ return list(self.agent_ids)
+
+ async def get_action_of_agent_without_side_effects(
+ self, agent_id: AgentId
+ ) -> Tuple[Any, AgentActLog]:
+ """
+ Safe function to get an action of an agent without modifying the agent or the simulation.
+ """
+ agent = self.agents[agent_id]
+ agent_before_action = agent.get_safe_copy()
+ obs = self.simulation.get_obs_agent(agent_id)
+ action, action_info = await agent.act(observation=obs)
+ self.agents[agent_id] = agent_before_action
+ agent_after_action = agent.get_safe_copy()
+ return AgentAndActionSafeCopy(action, action_info, agent_after_action)
+
+ async def get_actions_of_agents_without_side_effects(
+ self,
+ ) -> dict[AgentId, AgentAndActionSafeCopy]:
+ """
+ Safe function to get an action of an agent without modifying the agent or the simulation.
+ """
+ tasks = []
+ for agent_id in self.agent_ids:
+ task = asyncio.create_task(
+ self.get_action_of_agent_without_side_effects(agent_id)
+ )
+ tasks.append(task)
+ agent_and_action_safe_copies: list[
+ AgentAndActionSafeCopy
+ ] = await asyncio.gather(*tasks)
+ return {
+ agent_id: agent_and_action_safe_copy
+ for agent_id, agent_and_action_safe_copy in zip(
+ self.agent_ids, agent_and_action_safe_copies
+ )
+ }
+
+ def set_action_and_agent_after_action_manually(
+ self,
+ agent_id: AgentId,
+ agent_action_safe_copy: AgentAndActionSafeCopy,
+ ):
+ """
+ Set the action and the agent after action manually.
+ """
+ self.actions[agent_id] = agent_action_safe_copy.action
+ self.agent_step_logs[agent_id] = agent_action_safe_copy.action_info
+ self.agents[agent_id] = agent_action_safe_copy.agent_after_action
+
+ def set_actions_of_agents_manually(
+ self, actions: dict[AgentId, AgentAndActionSafeCopy]
+ ):
+ """
+ Set the actions of agents manually.
+ """
+ for agent_id, agent_action_safe_copy in actions.items():
+ self.set_action_and_agent_after_action_manually(
+ agent_id, agent_action_safe_copy
+ )
+
+ async def set_action_of_agent(self, agent_id: AgentId):
+ """
+ TOWRITE
+ """
+ agent = self.agents[agent_id]
+ obs = self.simulation.get_obs_agent(agent_id)
+ action, action_info = await agent.act(observation=obs)
+ self.actions[agent_id] = action
+ self.agent_step_logs[agent_id] = action_info
+
+ async def set_actions(self):
+ """
+ TOWRITE
+ """
+ # background_tasks = set()
+ tasks = []
+ for agent_id in self.agent_ids:
+ task = asyncio.create_task(self.set_action_of_agent(agent_id))
+ tasks.append(task)
+ await asyncio.gather(*tasks)
+
+ def take_simulation_step(self):
+ """
+ TOWRITE
+ """
+ terminated, self.simulation_step_log = self.simulation.step(self.actions)
+ return terminated
+
+ def get_step_log(self) -> StepLog:
+ """
+ TOWRITE
+ TODO: assert actions and simulation have taken step
+ """
+ step_log = StepLog(
+ simulation_step_log=self.simulation_step_log,
+ action_logs=self.agent_step_logs,
+ )
+ return step_log
+
+ async def step(self) -> Tuple[bool, StepLog]:
+ """
+ TOWRITE
+ """
+ await self.set_actions()
+ terminated = self.take_simulation_step()
+ step_log = self.get_step_log()
+ return terminated, step_log
+
+ def get_safe_copy(self):
+ """
+ TOWRITE
+ """
+
+ new_markov_game = copy.copy(self)
+ new_simulation = self.simulation.get_safe_copy()
+ new_agents = {
+ agent_id: agent.get_safe_copy() for agent_id, agent in self.agents.items()
+ }
+
+ # Reassign copied components
+ new_markov_game.simulation = new_simulation
+ new_markov_game.agents = new_agents
+
+ # IMPORTANT: ensure agent_ids references the new agents dict, not the original
+ new_markov_game.agent_ids = new_markov_game.agents.keys()
+
+ # Deep-copy step data to avoid correlation
+ new_markov_game.simulation_step_log = copy.deepcopy(self.simulation_step_log)
+ new_markov_game.actions = copy.deepcopy(self.actions)
+ # Rebuild logs to align exactly with new agent ids
+ old_agent_step_logs = copy.deepcopy(self.agent_step_logs)
+ new_markov_game.agent_step_logs = {
+ agent_id: old_agent_step_logs.get(agent_id)
+ for agent_id in new_markov_game.agent_ids
+ }
+
+ return new_markov_game
diff --git a/src_code_for_reproducibility/markov_games/mg_utils.py b/src_code_for_reproducibility/markov_games/mg_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..27e0eb606ea6cfd5682c3b21a53e783c23e9b43f
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/mg_utils.py
@@ -0,0 +1,89 @@
+import asyncio
+import copy
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from mllm.markov_games.ipd.ipd_agent import IPDAgent
+from mllm.markov_games.ipd.ipd_simulation import IPD
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.negotiation.dond_agent import DealNoDealAgent
+from mllm.markov_games.negotiation.dond_simulation import DealNoDealSimulation
+from mllm.markov_games.negotiation.nego_hard_coded_policies import (
+ HardCodedNegoGreedyPolicy,
+ HardCodedNegoWelfareMaximizingPolicy,
+)
+from mllm.markov_games.ipd.Ipd_hard_coded_agents import AlwaysCooperateIPDAgent, AlwaysDefectIPDAgent
+from mllm.markov_games.negotiation.no_press_nego_agent import NoPressAgent
+from mllm.markov_games.negotiation.no_press_nego_simulation import NoPressSimulation
+from mllm.markov_games.negotiation.tas_agent import TrustAndSplitAgent
+from mllm.markov_games.negotiation.tas_rps_agent import TrustAndSplitRPSAgent
+from mllm.markov_games.negotiation.tas_rps_simulation import TrustAndSplitRPSSimulation
+from mllm.markov_games.negotiation.tas_simple_agent import TrustAndSplitSimpleAgent
+from mllm.markov_games.negotiation.tas_simple_simulation import (
+ TrustAndSplitSimpleSimulation,
+)
+from mllm.markov_games.negotiation.tas_simulation import TrustAndSplitSimulation
+from mllm.markov_games.rollout_tree import (
+ AgentActLog,
+ RolloutTreeBranchNode,
+ RolloutTreeNode,
+ RolloutTreeRootNode,
+ StepLog,
+)
+from mllm.markov_games.simulation import SimulationStepLog
+
+AgentId = str
+
+
+@dataclass
+class AgentConfig:
+ agent_id: str
+ agent_name: str
+ agent_class_name: str
+ policy_id: str
+ init_kwargs: dict
+
+
+@dataclass
+class MarkovGameConfig:
+ id: int
+ seed: int
+ simulation_class_name: str
+ simulation_init_args: dict
+ agent_configs: list[AgentConfig]
+
+
+def init_markov_game_components(
+ config: MarkovGameConfig, policies: dict[str, Callable[[list[dict]], str]]
+):
+ """
+ TOWRITE
+ """
+ agents = {}
+ agent_names = []
+ for agent_config in config.agent_configs:
+ agent_id = agent_config.agent_id
+ agent_name = agent_config.agent_name
+ agent_class = eval(agent_config.agent_class_name)
+ agent = agent_class(
+ seed=config.seed,
+ agent_id=agent_id,
+ agent_name=agent_name,
+ policy=policies[agent_config.policy_id],
+ **agent_config.init_kwargs,
+ )
+ agents[agent_id] = agent
+ agent_names.append(agent_name)
+ simulation = eval(config.simulation_class_name)(
+ seed=config.seed,
+ agent_ids=list(agents.keys()),
+ agent_names=agent_names,
+ **config.simulation_init_args,
+ )
+ markov_game = MarkovGame(
+ id=config.id,
+ crn_id=config.seed,
+ agents=agents,
+ simulation=simulation,
+ )
+ return markov_game
diff --git a/src_code_for_reproducibility/markov_games/negotiation/nego_agent.py b/src_code_for_reproducibility/markov_games/negotiation/nego_agent.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b5bf4e3ca4ee7faa982360674e19d9eff6980dc
--- /dev/null
+++ b/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_str} 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=[]
+ )
diff --git a/src_code_for_reproducibility/markov_games/negotiation/nego_simulation.py b/src_code_for_reproducibility/markov_games/negotiation/nego_simulation.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a4d18532ab0472cd8f83414d52cf6df589fe126
--- /dev/null
+++ b/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
diff --git a/src_code_for_reproducibility/markov_games/negotiation/negotiation_statistics.py b/src_code_for_reproducibility/markov_games/negotiation/negotiation_statistics.py
new file mode 100644
index 0000000000000000000000000000000000000000..ccc7d53357033710b8409bdf2bfafafa58a40826
--- /dev/null
+++ b/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,
+]
diff --git a/src_code_for_reproducibility/markov_games/negotiation/no_press_nego_simulation.py b/src_code_for_reproducibility/markov_games/negotiation/no_press_nego_simulation.py
new file mode 100644
index 0000000000000000000000000000000000000000..d182187cc72c889a76f2d1c5be4b3afb6b923ed8
--- /dev/null
+++ b/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()
diff --git a/src_code_for_reproducibility/markov_games/negotiation/tas_rps_simulation.py b/src_code_for_reproducibility/markov_games/negotiation/tas_rps_simulation.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a4289f89a0574056024cdd5da0f8a676d331670
--- /dev/null
+++ b/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()
diff --git a/src_code_for_reproducibility/markov_games/negotiation/tas_simulation.py b/src_code_for_reproducibility/markov_games/negotiation/tas_simulation.py
new file mode 100644
index 0000000000000000000000000000000000000000..5499a146e9da491757a8105965b2d210f8327134
--- /dev/null
+++ b/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()
diff --git a/src_code_for_reproducibility/markov_games/rollout_tree.py b/src_code_for_reproducibility/markov_games/rollout_tree.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e90fa6d7a0e0b644ccaf3533a1513568d21868f
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/rollout_tree.py
@@ -0,0 +1,86 @@
+"""
+TODO: add parent to nodes so that some verification can be done. For instance, to ensure that node reward keys match the parent node.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, List, Literal, Optional, Tuple
+
+import jsonschema
+from pydantic import BaseModel, Field, model_validator
+
+from mllm.chat_utils.chat_turn import ChatTurn
+
+AgentId = str
+
+
+class SimulationStepLog(BaseModel):
+ rewards: dict[AgentId, float]
+ info: Any = None
+
+
+class AgentActLog(BaseModel):
+ chat_turns: list[ChatTurn] | None
+ info: Any = None
+
+ @model_validator(mode="after")
+ def _exactly_one_state_end(self):
+ """
+ This method is used to enforce that for each AgentActLog, there is exactly one ChatTurn which is a state end.
+ """
+ if self.chat_turns != []:
+ n = sum(1 for t in self.chat_turns if t.is_state_end)
+ if n != 1:
+ raise ValueError(
+ f"AgentActLog must have exactly one ChatTurn with is_state_end=True; got {self.chat_turns}."
+ )
+ return self
+ else:
+ return self
+
+
+class StepLog(BaseModel):
+ action_logs: dict[AgentId, AgentActLog]
+ simulation_step_log: SimulationStepLog
+
+
+# BranchType = Literal["unilateral_deviation", "common_deviation"] # might not be necessary
+# class BranchNodeInfo(BaseModel):
+# branch_id: str
+# branch_for: AgentId
+# branch_type: BranchType
+
+
+class RolloutTreeNode(BaseModel):
+ step_log: StepLog
+ time_step: int
+ child: RolloutTreeNode | RolloutTreeBranchNode | None = None
+
+
+class RolloutTreeBranchNode(BaseModel):
+ """
+ First item of the tuple indicates which agent "called" for an alternative branch.
+ """
+
+ main_child: RolloutTreeNode
+ branches: dict[AgentId, list[RolloutTreeNode]] | None = None
+
+
+class RolloutTreeRootNode(BaseModel):
+ id: int
+ crn_id: int # ID of the rng used to generate this rollout tree
+ child: RolloutTreeNode | RolloutTreeBranchNode | None = None
+ agent_ids: List[AgentId] = Field(min_length=1)
+
+
+# class RolloutTreeLeafNode(BaseModel):
+# step_log: StepLog
+# time_step: int
+
+
+# Necessary for self-referential stuff in pydantic
+RolloutTreeBranchNode.model_rebuild()
+RolloutTreeNode.model_rebuild()
diff --git a/src_code_for_reproducibility/markov_games/run_markov_games.py b/src_code_for_reproducibility/markov_games/run_markov_games.py
new file mode 100644
index 0000000000000000000000000000000000000000..08b84024668ed4375453d1fd515a78eb9bb23414
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/run_markov_games.py
@@ -0,0 +1,24 @@
+import asyncio
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from torch._C import ClassType
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import RolloutTreeRootNode
+
+
+async def run_markov_games(
+ runner: Callable[[MarkovGame], RolloutTreeRootNode],
+ runner_kwargs: dict,
+ output_folder: str,
+ markov_games: list[MarkovGame],
+) -> list[RolloutTreeRootNode]:
+ tasks = []
+ for mg in markov_games:
+ tasks.append(
+ asyncio.create_task(
+ runner(markov_game=mg, output_folder=output_folder, **runner_kwargs)
+ )
+ )
+ return await asyncio.gather(*tasks)
diff --git a/src_code_for_reproducibility/markov_games/simulation.py b/src_code_for_reproducibility/markov_games/simulation.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0b804e2aa4c288b3d98cc8106cfd727f1cc1e1a
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/simulation.py
@@ -0,0 +1,87 @@
+"""
+A Simulation is the environment of a Markov Game.
+The Simulation is not responsible for properly checking / formatting the responses of LLM's.
+This is the job of the `Agent` class.
+Simulations expect clean actions, and are defined similarly to `gymnasium` environments, except that they are adapted for the Multi-agent setting.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Tuple
+
+from numpy.random import default_rng
+
+from mllm.markov_games.rollout_tree import SimulationStepLog
+
+
+class Simulation(ABC):
+ @abstractmethod
+ def __init__(self, seed: int, *args, **kwargs):
+ self.seed = seed
+ self.rng = default_rng(self.seed)
+
+ @abstractmethod
+ def step(self, actions: Any) -> Tuple[bool, SimulationStepLog]:
+ """
+ Returns terminated, info
+ """
+ raise NotImplementedError
+
+ def get_obs(self):
+ """Returns all agent observations in dict
+
+ Returns:
+ observations
+ """
+ raise NotImplementedError
+
+ def get_obs_agent(self, agent_id):
+ """Returns observation for agent_id"""
+ raise NotImplementedError
+
+ def get_obs_size(self):
+ """Returns the shape of the observation"""
+ raise NotImplementedError
+
+ def get_state(self):
+ raise NotImplementedError
+
+ def get_state_size(self):
+ """Returns the shape of the state"""
+ raise NotImplementedError
+
+ def get_avail_actions(self):
+ raise NotImplementedError
+
+ def get_avail_agent_actions(self, agent_id):
+ """Returns the available actions for agent_id"""
+ raise NotImplementedError
+
+ def get_total_actions(self):
+ """Returns the total number of actions an agent could ever take"""
+ # TODO: This is only suitable for a discrete 1 dimensional action space for each agent
+ raise NotImplementedError
+
+ def get_safe_copy(self):
+ """
+ Return copy of the agent object that is decorrelated from the original object.
+ """
+ raise NotImplementedError
+
+ def reset(self):
+ """Returns initial observations and states"""
+ raise NotImplementedError
+
+ def render(self):
+ raise NotImplementedError
+
+ def close(self):
+ raise NotImplementedError
+
+ # def seed(self):
+ # raise NotImplementedError
+
+ def save_replay(self):
+ raise NotImplementedError
+
+ def get_simulation_info(self):
+ raise NotImplementedError
diff --git a/src_code_for_reproducibility/markov_games/vine_ppo.py b/src_code_for_reproducibility/markov_games/vine_ppo.py
new file mode 100644
index 0000000000000000000000000000000000000000..f3cd2c89133fa00b0d1d0fa260688efa642d4ded
--- /dev/null
+++ b/src_code_for_reproducibility/markov_games/vine_ppo.py
@@ -0,0 +1,10 @@
+from anytree import Node, RenderTree
+from anytree.exporter import DotExporter
+import os.path
+import asyncio
+from mllm.markov_games.markov_game import MarkovGame
+
+async def VinePPORunner(
+ markov_game: MarkovGame,
+ **kwargs):
+ pass
diff --git a/src_code_for_reproducibility/models/__pycache__/__init__.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..273d114b579eaaff7311e14e8c9746b08c9093b5
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/__init__.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/adapter_training_wrapper.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/adapter_training_wrapper.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..66b8c4e783fd86c78c01fb1fe2de031ffa1363fc
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/adapter_training_wrapper.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/human_policy.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/human_policy.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e3f4a2685a9cd79a8828cb18f6a3d32f03e63769
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/human_policy.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/inference_backend.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/inference_backend.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c77324d067890d56a074a4e427cb5e9b6c59df01
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/inference_backend.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/inference_backend_dummy.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/inference_backend_dummy.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6a89b2fde9bcdd46de8e225fa9ee9e67ae9f5b9c
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/inference_backend_dummy.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/inference_backend_sglang.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/inference_backend_sglang.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..84ec9701c778683849c6479de32b44b86a1ec82b
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/inference_backend_sglang.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/inference_backend_vllm.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/inference_backend_vllm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..da0a8fd5d1dceb7ea6f095412c40fd34b8534e9b
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/inference_backend_vllm.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/large_language_model_api.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/large_language_model_api.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7de6daa0ac310a269a94692e05457e6c5a8ab925
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/large_language_model_api.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/large_language_model_local.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/large_language_model_local.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e1f8ad83ac621376945ef4f69083606ae9d8bca4
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/large_language_model_local.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/models/__pycache__/scalar_critic.cpython-312.pyc b/src_code_for_reproducibility/models/__pycache__/scalar_critic.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..650e0c11acbe423288c4a4a4cd005fe2be810eea
Binary files /dev/null and b/src_code_for_reproducibility/models/__pycache__/scalar_critic.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/__init__.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6d68eb043e149c6005bbcdfe2fbfb1e1d1d61819
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/__init__.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/annealing_methods.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/annealing_methods.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..73991155f0a6d13f03b43dfcd777a49610c88c0e
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/annealing_methods.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/credit_methods.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/credit_methods.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..549ee7cccb7bf669178129df112ae091e2389518
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/credit_methods.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/produce_training_stats.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/produce_training_stats.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..eeb7e4551d233ca00f0346838c87e7abbdb84fa1
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/produce_training_stats.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/tally_basic.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tally_basic.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..96098409cb3e042778965df4c94f7b2a45a66101
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tally_basic.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/tally_metrics.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tally_metrics.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6c2ee27df0f569f1b270d0b7e2977afa9c954181
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tally_metrics.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/tally_rollout.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tally_rollout.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c27a1b8c25697b30b7642894b4ac32e5f5b41531
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tally_rollout.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/tally_tokenwise.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tally_tokenwise.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..afe36bd9eb6fa41bb714296a853748846986ed3c
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tally_tokenwise.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/tokenize_chats.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/tokenize_chats.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bb3cd1234f6a70b27bef4626f66ec30d5ab52517
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/tokenize_chats.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/trainer_ad_align.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/trainer_ad_align.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7a1cc78df9edd45b4e52fc1f2685eb3a99d48d23
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/trainer_ad_align.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/trainer_common.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/trainer_common.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cc7e9a6c2510ae34731aa3d46ca42829a483bba4
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/trainer_common.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/trainer_independent.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/trainer_independent.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..45d61bbfffaa2299ef7c915156ead08105ef64af
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/trainer_independent.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/trainer_sum_rewards.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/trainer_sum_rewards.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ad600757f4ff3da914a94f33768dba90070148cb
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/trainer_sum_rewards.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/training/__pycache__/training_data_utils.cpython-312.pyc b/src_code_for_reproducibility/training/__pycache__/training_data_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..01210498226a9a10c4c6675edd9108c376cad74b
Binary files /dev/null and b/src_code_for_reproducibility/training/__pycache__/training_data_utils.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/__init__.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..32c4e0f8b0b09f12cba13137dbdc2a6f1d367f0f
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/__init__.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/dict_get_path.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/dict_get_path.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ec38a350f3ab2827c91067693e5aaed5cd6e1385
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/dict_get_path.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/gather_training_stats.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/gather_training_stats.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..09bf572c7c52dd7aab1bb6626ad794e15a4b2c55
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/gather_training_stats.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/get_coagent_id.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/get_coagent_id.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..05691ab328cc6ccf7a54fc828757fc12582c8ec5
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/get_coagent_id.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/kill_sglang.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/kill_sglang.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..98167c99cd89bba3ea9b6cb75d28572034a683c9
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/kill_sglang.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/resource_context.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/resource_context.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6b9763f1e29c04e7debd256efe7bcfc0cdcfefc8
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/resource_context.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/rollout_tree_chat_htmls.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_chat_htmls.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..33285a6e93f9825a14855114ec1c49eeb89abd74
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_chat_htmls.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/rollout_tree_gather_utils.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_gather_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fdd056e5ba1aceb1536f2eab49f1c63fe7fd647d
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_gather_utils.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/rollout_tree_stats.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_stats.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5eb56226df7b241e081728dc406c4f0162cba1fa
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/rollout_tree_stats.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/short_id_gen.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/short_id_gen.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..10898f798ab2178b950e57100739d421c01f00f7
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/short_id_gen.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/stat_pack.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/stat_pack.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..65e397d163b96f1a0b90c7541734bd729bb82c9e
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/stat_pack.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/update_start_epoch.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/update_start_epoch.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3ed6ba62682410f9b2fe69fed3538a96d8fda887
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/update_start_epoch.cpython-312.pyc differ
diff --git a/src_code_for_reproducibility/utils/__pycache__/wandb_utils.cpython-312.pyc b/src_code_for_reproducibility/utils/__pycache__/wandb_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..09bd9177e48a4140805ea2695cb567ffe987d70a
Binary files /dev/null and b/src_code_for_reproducibility/utils/__pycache__/wandb_utils.cpython-312.pyc differ