Add files using upload-large-folder tool

src_code_for_reproducibility/docs/source/conf.py

@@ -0,0 +1,48 @@
+# Configuration file for the Sphinx documentation builder.
+import os
+import sys
+sys.path.insert(0, os.path.abspath('../..'))
+
+# -- Project information -----------------------------------------------------
+project = 'llm_negotiation'
+copyright = '2023, Your Name'
+author = 'Your Name'
+
+# -- General configuration ---------------------------------------------------
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.viewcode',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.autosummary',
+    'sphinx.ext.intersphinx',
+    'sphinx.ext.mathjax',
+    'sphinxcontrib.mermaid',
+    'sphinx_rtd_theme',
+]
+
+templates_path = ['_templates']
+exclude_patterns = []
+
+# -- Options for HTML output -------------------------------------------------
+html_theme = 'sphinx_rtd_theme'
+html_static_path = ['_static']
+
+# -- Napoleon settings -------------------------------------------------------
+napoleon_google_docstring = True
+napoleon_numpy_docstring = False
+napoleon_include_init_with_doc = True
+napoleon_include_private_with_doc = False
+napoleon_include_special_with_doc = True
+napoleon_use_admonition_for_examples = False
+napoleon_use_admonition_for_notes = False
+napoleon_use_admonition_for_references = False
+napoleon_use_ivar = False
+napoleon_use_param = True
+napoleon_use_rtype = True
+napoleon_preprocess_types = False
+napoleon_type_aliases = None
+napoleon_attr_annotations = True
+
+# -- Path setup --------------------------------------------------------------
+# Make sure the project's modules can be found by Sphinx
+sys.path.insert(0, os.path.abspath('../../src'))

src_code_for_reproducibility/docs/source/environments/diplomacy.rst

@@ -0,0 +1,459 @@
+=================
+Diplomacy
+=================
+
+The Diplomacy environment provides a multi-agent negotiation interface for the classic board game Diplomacy, 
+based on DeepMind's implementation. This document describes the API for interacting with the Diplomacy environment
+and its associated agent handler.
+
+Overview
+--------
+
+Diplomacy is a strategic board game set in Europe before World War I, where players control one of seven European powers 
+and negotiate with each other to gain control of supply centers. The game is played in turns, with each turn consisting 
+of movement phases, retreat phases, and build phases.
+
+Our implementation adapts DeepMind's Diplomacy code to the Multi-Agent Negotiation Environment standard, allowing it 
+to be used with LLM agents through a text-based interface.
+
+Game Rules
+----------
+
+### Game Board and Powers
+
+Diplomacy is played on a map of Europe divided into provinces. The game features seven Great Powers that players can control:
+
+- England (blue)
+- France (light blue)
+- Germany (black)
+- Italy (green)
+- Austria-Hungary (red)
+- Russia (white)
+- Turkey (yellow)
+
+Each power begins with three supply centers (except Russia, which starts with four) and an equal number of units.
+
+### Units and Movement
+
+There are two types of units in Diplomacy:
+- **Armies (A)**: Can move to adjacent land provinces or be convoyed across water by fleets
+- **Fleets (F)**: Can move to adjacent coastal provinces and sea regions
+
+During movement phases, each unit can execute one of these orders:
+- **Hold**: The unit remains in its current province (e.g., "A PAR H")
+  - Format: [Unit Type] [Province] H
+  - Example: "A PAR H" means "Army in Paris holds its position"
+
+- **Move**: The unit attempts to move to an adjacent province (e.g., "A PAR - BUR")
+  - Format: [Unit Type] [Current Province] - [Destination Province]
+  - Example: "A PAR - BUR" means "Army in Paris moves to Burgundy"
+  - Example: "F BRE - ENG" means "Fleet in Brest moves to the English Channel"
+
+- **Support**: The unit supports another unit's move or hold (e.g., "A PAR S A MAR - BUR")
+  - Format for supporting a move: [Unit Type] [Province] S [Unit Type] [Province] - [Destination]
+  - Format for supporting a hold: [Unit Type] [Province] S [Unit Type] [Province]
+  - Example: "A PAR S A MAR - BUR" means "Army in Paris supports the Army in Marseille's move to Burgundy"
+  - Example: "F LON S F NTH" means "Fleet in London supports the Fleet in North Sea holding its position"
+
+- **Convoy**: A fleet can convoy an army across water (e.g., "F ENG C A LON - BRE")
+  - Format: [Fleet] [Sea Province] C [Army] [Coastal Province] - [Coastal Province]
+  - Example: "F ENG C A LON - BRE" means "Fleet in English Channel convoys the Army in London to Brest"
+
+All orders are executed simultaneously, and conflicts are resolved based on strength (number of supporting units).
+
+### Common Province Abbreviations
+
+Diplomacy uses three-letter abbreviations for provinces. Some common ones include:
+- **PAR**: Paris
+- **LON**: London
+- **BER**: Berlin
+- **MUN**: Munich
+- **BUR**: Burgundy
+- **MAR**: Marseilles
+- **BRE**: Brest
+- **ENG**: English Channel
+- **NTH**: North Sea
+- **VIE**: Vienna
+- **ROM**: Rome
+- **VEN**: Venice
+- **MOW**: Moscow
+- **CON**: Constantinople
+
+### Example: Movement and Conflicts
+
+For example, if France orders "A PAR - BUR" and Germany orders "A MUN - BUR", neither move succeeds as they have equal strength. However, if France also orders "A MAR S A PAR - BUR", then the French army from Paris would successfully move to Burgundy with strength of 2 against Germany's strength of 1.
+
+### Turn Structure
+
+A game year consists of five phases:
+1. **Spring Movement**: All powers submit orders for their units
+2. **Spring Retreat**: Units dislodged in the movement phase must retreat or be disbanded
+3. **Fall Movement**: Another round of movement orders
+4. **Fall Retreat**: Retreat orders for dislodged units
+5. **Winter Adjustment**: Powers gain or lose units based on the number of supply centers they control
+
+### Supply Centers and Building
+
+Supply centers (marked on the map) are key to victory. When a power occupies a supply center during a Fall turn, they gain control of it. During the Winter Adjustment phase:
+- If you control more supply centers than you have units, you can build new units in your home supply centers
+- If you control fewer supply centers than you have units, you must remove excess units
+
+### Example: Building and Removing Units
+
+If France controls 5 supply centers but only has 4 units, during the Winter phase they can build one new unit in an unoccupied home supply center (Paris, Marseilles, or Brest). Conversely, if France controls only 3 supply centers but has 4 units, they must remove one unit of their choice.
+
+### Negotiation
+
+A critical component of Diplomacy is the negotiation between players. Before submitting orders, players can communicate freely to form alliances, coordinate attacks, or mislead opponents. These negotiations are not binding, and betrayal is a common strategy.
+
+### Example: Alliance and Betrayal
+
+England and France might agree to an alliance against Germany, with England promising to support France's move into Belgium. However, England could secretly order their fleet to move into Belgium themselves or support a German move instead.
+
+### Victory Conditions
+
+The game ends when one power controls 18 or more supply centers (majority of the 34 total centers), or when players agree to a draw. In tournament settings, games may also end after a predetermined number of game years.
+
+DiplomacyEnv
+------------
+
+The ``DiplomacyEnv`` class provides an interface to the Diplomacy game environment that follows the Multi-Agent 
+Negotiation Environment standard.
+
+.. code-block:: python
+
+    class DiplomacyEnv:
+        """
+        Multi-Agent Negotiation Environment for Diplomacy, adapting Deepmind's implementation
+        to the MarlEnvironment standard.
+        """
+        def __init__(self, 
+                    initial_state: Optional[DiplomacyState] = None,
+                    max_turns: int = 100,
+                    points_per_supply_centre: bool = True,
+                    forced_draw_probability: float = 0.0,
+                    min_years_forced_draw: int = 35):
+            """Initialize the Diplomacy environment.
+            
+            Args:
+                initial_state: Initial DiplomacyState (optional)
+                max_turns: Maximum number of turns in the game
+                points_per_supply_centre: Whether to award points per supply center in case of a draw
+                forced_draw_probability: Probability of forcing a draw after min_years_forced_draw
+                min_years_forced_draw: Minimum years before considering a forced draw
+            """
+            # ...
+            
+        def reset(self):
+            """Reset the environment to an initial state and return the initial observation.
+            
+            Returns:
+                observation (dict): A dictionary where keys are agent identifiers and values are observations.
+                Each observation contains:
+                - board_state: Current state of the board
+                - current_season: Current season in the game
+                - player_index: Index of the player's power
+                - possible_actions: List of possible actions in DeepMind's format
+                - human_readable_actions: List of human-readable action descriptions
+                - supply_centers: List of supply centers owned by the player
+                - units: List of units owned by the player
+                - year: Current year in the game
+            """
+            # ...
+            
+        def step(self, actions):
+            """Take a step in the environment using the provided actions.
+
+            Args:
+                actions (dict): A dictionary where keys are agent identifiers and values are actions.
+                    Actions can be:
+                    - List of integer actions in DeepMind's format
+                    - List of string actions in text format (e.g., "A MUN - BER")
+
+            Returns:
+                observations (dict): A dictionary where keys are agent identifiers and values are observations.
+                    Each observation has the same structure as in reset().
+                done (bool): Whether the episode has ended.
+                info (dict): Additional information about the environment, including:
+                    - turn: Current turn number
+                    - returns: Game returns if the game is done, otherwise None
+                    - waiting_for: List of agents that still need to provide actions (if not all actions are provided)
+            """
+            # ...
+            
+        def get_log_info(self):
+            """Get additional information about the environment for logging.
+            
+            Returns:
+                log_info (dict): Information about the environment required to log the game, including:
+                    - power_names: List of power names
+                    - game_history: History of the game
+                    - current_turn: Current turn number
+                    - current_season: Current season name
+                    - supply_centers: Dictionary mapping power names to supply center counts
+            """
+            # ...
+            
+        def render(self):
+            """Render the current state of the environment.
+            
+            Displays a visualization of the current game state.
+            """
+            # ...
+            
+        def close(self):
+            """Perform any necessary cleanup."""
+            # ...
+
+
+Key Implementation Details
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The ``DiplomacyEnv`` class implements several key features:
+
+1. **Multi-Agent Support**: The environment tracks multiple agents (powers) and manages their interactions.
+
+2. **Turn-Based Gameplay**: The environment enforces the turn structure of Diplomacy, including different phases.
+
+3. **Action Processing**: The environment can handle actions in both text format and DeepMind's integer format.
+
+4. **Observation Generation**: The environment generates detailed observations for each agent, including board state, supply centers, and possible actions.
+
+5. **Game Termination**: The environment tracks game termination conditions, including supply center victory and maximum turn limits.
+
+Observation Structure
+~~~~~~~~~~~~~~~~~~~~
+
+Each agent receives an observation dictionary with the following structure:
+
+.. code-block:: python
+
+    {
+        "board_state": np.ndarray,  # Board state representation
+        "current_season": int,      # Season index (0-4)
+        "player_index": int,        # Index of the player's power (0-6)
+        "possible_actions": [int],  # List of possible actions in DeepMind's format
+        "human_readable_actions": [str],  # List of human-readable action descriptions
+        "supply_centers": [str],    # List of supply centers owned by the player
+        "units": [dict],            # List of units owned by the player
+        "year": int                 # Current year in the game
+    }
+
+Action Structure
+~~~~~~~~~~~~~~~
+
+Actions can be provided in two formats:
+
+1. **Text Format**: String actions like ``"A MUN - BER"`` or ``"F NTH C A LON - BEL"``.
+
+2. **Integer Format**: Lists of integers corresponding to DeepMind's action representation.
+
+The environment will convert text actions to the internal format as needed.
+
+DiplomacyAgent
+--------------
+
+The ``DiplomacyAgent`` class implements the agent handler interface for Diplomacy, processing observations from the environment and generating actions through an LLM.
+
+.. code-block:: python
+
+    class DiplomacyAgent:
+        """
+        Agent handler for Diplomacy, implementing the AgentState interface
+        for the multi-agent negotiation standard.
+        """
+        
+        def __init__(self, 
+                    power_name: str,
+                    use_text_interface: bool = True,
+                    system_prompt: Optional[str] = None):
+            """Initialize the Diplomacy agent handler.
+            
+            Args:
+                power_name: Name of the power this agent controls
+                use_text_interface: Whether to use text-based interface (vs. structured)
+                system_prompt: Optional system prompt to use for the LLM
+            """
+            # ...
+            
+        def step(self, observation_from_env, policy_output=None):
+            """Update the agent state based on the observation and action.
+            
+            Args:
+                observation_from_env: The observation from the environment, with structure:
+                    - board_state: Current state of the board
+                    - current_season: Current season in the game
+                    - player_index: Index of the player's power
+                    - possible_actions: List of possible actions
+                    - human_readable_actions: List of human-readable action descriptions
+                    - supply_centers: List of supply centers owned by the player
+                    - units: List of units owned by the player
+                    - year: Current year in the game
+                
+                policy_output: The output of the policy (LLM response), or None for initial prompt
+                
+            Returns:
+                policy_id (str): The policy identifier ("llm_policy")
+                policy_input (dict): The input to the policy, with structure:
+                    - messages: List of conversation messages in the format:
+                        [{"role": "system", "content": "..."}, 
+                         {"role": "user", "content": "..."}]
+                action: The official action to be sent to the environment, or None if not ready
+                done (bool): Whether the LLM action is ready to be sent to the environment
+                info (dict): Additional information about the agent:
+                    - valid_action: Whether the extracted action is valid
+            """
+            # ...
+            
+        def get_log_info(self):
+            """Get information about the agent required to log a trajectory.
+            
+            Returns:
+                log_info (dict): Information about the agent required to log a trajectory:
+                    - power_name: Name of the power this agent controls
+                    - conversation_history: List of conversation messages
+                    - current_action: The current action, if any
+            """
+            # ...
+            
+        def render(self):
+            """Render the current state of the agent.
+            
+            Displays the agent's current state, including conversation history.
+            """
+            # ...
+            
+        def close(self):
+            """Perform any necessary cleanup."""
+            # ...
+
+
+Key Implementation Details
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The ``DiplomacyAgent`` class implements several key features:
+
+1. **LLM Interaction**: The agent generates prompts for an LLM and processes the LLM's responses to extract actions.
+
+2. **Conversation Management**: The agent maintains a conversation history for coherent interactions with the LLM.
+
+3. **Action Validation**: The agent validates extracted actions against the set of possible actions provided by the environment.
+
+4. **Error Handling**: The agent generates clarification prompts when invalid actions are detected.
+
+5. **Text-Based Interface**: The agent formats game state information into human-readable text for the LLM.
+
+Prompt Structure
+~~~~~~~~~~~~~~~
+
+The agent generates prompts that include:
+
+1. **System Prompt**: Instructions and context for the LLM, explaining its role as a Diplomacy player.
+
+2. **Game State Description**: A text description of the current game state, including:
+   - Current year and season
+   - Supply centers owned
+   - Units controlled
+   - Possible actions
+
+3. **Action Request**: Instructions on how to format actions.
+
+Example system prompt:
+
+.. code-block:: text
+
+    You are playing the role of FRANCE in a game of Diplomacy. 
+    Your goal is to control as many supply centers as possible. 
+    You can negotiate with other players and form alliances, but remember that 
+    these alliances are not binding. When you need to submit orders for your units,
+    write them in the correct format, with each order on a new line.
+
+Example game state description:
+
+.. code-block:: text
+
+    Year: 1901, Season: SPRING_MOVES
+    You are playing as FRANCE.
+    You currently control 3 supply centers: PAR, MAR, BRE.
+    Your units are: A PAR, A MAR, F BRE.
+
+    Please provide orders for your units. Here are your possible actions:
+    A PAR - BUR
+    A PAR - GAS
+    A PAR - PIC
+    A PAR H
+    ...
+
+    Submit your orders, one per line, in the format like: "A MUN - BER" or "F NTH C A LON - BEL"
+
+Running Diplomacy Games
+----------------------
+
+To run Diplomacy games with LLM agents, you can use the ``run_batched_matches`` function with the ``DiplomacyEnv`` and ``DiplomacyAgent`` classes:
+
+.. code-block:: python
+
+    from mllm.environments.diplomacy.diplomacy_env import DiplomacyEnv
+    from mllm.environments.diplomacy.diplomacy_agent import DiplomacyAgent
+    from mllm.run_matches import run_batched_matches
+
+    # Create environment and agent handlers
+    env = DiplomacyEnv(max_turns=30)
+    
+    agent_handlers = {
+        "AUSTRIA": DiplomacyAgent(power_name="AUSTRIA"),
+        "ENGLAND": DiplomacyAgent(power_name="ENGLAND"),
+        "FRANCE": DiplomacyAgent(power_name="FRANCE"),
+        "GERMANY": DiplomacyAgent(power_name="GERMANY"),
+        "ITALY": DiplomacyAgent(power_name="ITALY"),
+        "RUSSIA": DiplomacyAgent(power_name="RUSSIA"),
+        "TURKEY": DiplomacyAgent(power_name="TURKEY")
+    }
+
+    # Define policy mapping (mapping from policy IDs to actual policy functions)
+    policy_mapping = {
+        "llm_policy": my_llm_policy_function
+    }
+
+    # Run the game
+    game_results = run_batched_matches(
+        envs=[env],
+        agent_handlers_per_env=[agent_handlers],
+        policy_mapping=policy_mapping,
+        max_parallel_matches=1
+    )
+
+    # Process results
+    for result in game_results:
+        print(f"Game finished. Winner: {result['winner']}")
+        print(f"Supply centers: {result['supply_centers']}")
+
+This setup allows you to run Diplomacy games with LLM agents using the Multi-Agent Negotiation Environment standard.
+
+Limitations and Considerations
+-----------------------------
+
+1. **Performance**: Processing observations and actions for seven powers using LLMs can be computationally intensive.
+
+2. **Action Parsing**: Extracting valid actions from LLM outputs may require sophisticated parsing and error handling.
+
+3. **Game Complexity**: Diplomacy is a complex game with many rules and edge cases, which may be challenging for LLMs to fully grasp.
+
+4. **Turn Duration**: Real Diplomacy games include negotiation phases of variable duration, which are not fully captured in this implementation.
+
+5. **Text Formatting**: The quality of LLM interactions depends heavily on the formatting and clarity of text prompts.
+
+Advanced Usage
+------------
+
+For advanced usage, you can customize:
+
+1. **System Prompts**: Modify agent behavior by providing custom system prompts.
+
+2. **Observation Processing**: Extend the observation processing to include additional information.
+
+3. **Action Parsing**: Implement more sophisticated action parsing for complex orders.
+
+4. **Visualization**: Add custom visualization methods to the environment's render function.
+
+5. **Logging**: Extend the logging capabilities to capture additional information about the game state. 

src_code_for_reproducibility/docs/source/environments/dond.rst

@@ -0,0 +1,410 @@
+=================
+Deal or No Deal
+=================
+
+The Deal or No Deal (DoND) environment provides a multi-agent negotiation interface where players trade 
+items with different values. This document describes the API for interacting with the DoND environment
+and its associated agent handler.
+
+Overview
+--------
+
+Deal or No Deal is a negotiation game where two agents must agree on how to divide a set of items, 
+each of which has different values to each agent. The agents engage in a back-and-forth dialogue to 
+determine an allocation of the items, with each trying to maximize their own total value.
+
+Our implementation follows the Multi-Agent Negotiation Environment standard, allowing it to be used 
+with LLM agents through a text-based interface.
+
+Game Rules
+----------
+
+### Basic Structure
+
+The core mechanics of Deal or No Deal are:
+
+1. Two agents negotiate over a set of items (e.g., books, balls, hats)
+2. Each item has:
+   - A specific quantity (how many of each item is available)
+   - A value for each agent (which may differ between agents)
+3. Agents take turns sending messages to negotiate how to split the items
+4. Once an agreement is reached, agents finalize the deal
+5. Points are awarded based on the value of items each agent receives
+
+### Detailed Gameplay
+
+#### Setup Phase
+
+The game begins with:
+- A set of items (e.g., "book", "hat", "ball")
+- Each item has a quantity (e.g., 6 books, 2 hats, 4 balls)
+- Each agent has private values for each item (e.g., books might be worth 5 points to one agent but only 2 points to the other)
+- Agents are assigned roles (starting negotiator and responding negotiator)
+
+#### Negotiation Phase
+
+1. Agents take turns sending free-form text messages to each other
+2. Messages can include offers, counter-offers, questions, or strategic communication
+3. There is a maximum number of messages permitted (preventing endless negotiations)
+4. Either agent can propose to finalize an agreement at any time
+
+For example:
+- Agent 1: "I propose I get all the books and you get all the hats and balls."
+- Agent 2: "That doesn't work for me. How about you get 3 books and I get 3 books, all the hats, and all the balls?"
+- Agent 1: "Let me counter-offer: I get 4 books and 2 balls, you get 2 books, all hats, and 2 balls."
+
+#### Finalization Phase
+
+1. When an agent wants to finalize a deal, they must specify the exact allocation:
+   - How many of each item they receive
+   - How many of each item the other agent receives
+2. The other agent must then either agree (by submitting the same allocation) or reject the finalization
+3. If both agents submit matching finalizations, the deal is executed
+4. If finalizations don't match, no agreement is reached, and both agents receive 0 points
+
+#### Scoring
+
+1. Each agent's score is calculated based on the value of items they receive
+2. The formula is: Sum(quantity_of_item_i × value_of_item_i_to_agent)
+3. If no agreement is reached, both agents receive 0 points
+
+### Example Game
+
+Let's walk through a simple example:
+
+**Setup:**
+- Items: Books (4), Hats (2), Balls (6)
+- Agent 1 values: Books=5, Hats=1, Balls=2
+- Agent 2 values: Books=3, Hats=6, Balls=1
+
+**Negotiation (simplified):**
+1. Agent 1: "I would like all the books and balls. You can have the hats."
+2. Agent 2: "That doesn't work for me. Books are valuable. I propose I get all the hats and 2 books, you get 2 books and all the balls."
+3. Agent 1: "How about I get 3 books and all the balls, and you get 1 book and all the hats?"
+4. Agent 2: "I accept your proposal."
+
+**Finalization:**
+- Agent 1 submits: Agent 1 gets (Books: 3, Hats: 0, Balls: 6), Agent 2 gets (Books: 1, Hats: 2, Balls: 0)
+- Agent 2 submits the same allocation, confirming agreement
+
+**Scoring:**
+- Agent 1 score: (3 books × 5) + (0 hats × 1) + (6 balls × 2) = 15 + 0 + 12 = 27 points
+- Agent 2 score: (1 book × 3) + (2 hats × 6) + (0 balls × 1) = 3 + 12 + 0 = 15 points
+
+### Game Variations
+
+The DoND environment supports several variations through configuration parameters:
+
+#### Different Value Distributions
+
+The environment offers multiple ways to assign values to items:
+
+1. **Standard Random Setup (dond_random_setup)**:
+   - Items have even-numbered quantities
+   - Each agent receives distinct random values for each item
+   - Values are drawn from a uniform distribution
+
+2. **Independent Random Values (independent_random_vals)**:
+   - Item quantities can be any number in the specified range
+   - Values for each agent are drawn independently
+   - Creates more varied negotiation scenarios
+
+3. **Bicameral Value Distribution (bicameral_vals_assignator)**:
+   - Creates a "high value" and "low value" distribution for each item
+   - Each agent values approximately half the items highly and half lowly
+   - Values are drawn from normal distributions with different means
+   - Creates scenarios with clear trade opportunities
+
+#### Visibility Options
+
+1. **Finalization Visibility**:
+   - When enabled, both agents can see each other's finalization proposals
+   - When disabled, finalization proposals remain private until both are submitted
+
+2. **Other Values Visibility**:
+   - When enabled, agents can see each other's value functions
+   - When disabled, agents only know their own values
+   - Creates information asymmetry and richer negotiation dynamics
+
+#### Game Modes
+
+1. **Cooperative Mode ("coop")**:
+   - Agents are encouraged to find mutually beneficial solutions
+   - Success is measured by the sum of both agents' scores
+
+2. **Competitive Mode ("comp")**:
+   - Agents aim to maximize their individual scores
+   - Creates more adversarial negotiations
+
+#### Round Structure
+
+1. **Single Round**:
+   - One negotiation session between the same agents
+   - Simple evaluation of negotiation skills
+
+2. **Multiple Rounds**:
+   - Agents negotiate multiple times with different item setups
+   - Allows for learning and adaptation over time
+   - Roles can be swapped between rounds
+
+DondEnv
+------------
+
+The ``DondEnv`` class provides an interface to the Deal or No Deal environment that follows the Multi-Agent 
+Negotiation Environment standard.
+
+.. code-block:: python
+
+    class DondEnv:
+        """
+        Multi-Agent Negotiation Environment for Deal or No Deal.
+        """
+        def __init__(
+            self,
+            agents,
+            mode="coop",
+            max_messages=None,
+            min_messages=None,
+            max_chars_per_message=None,
+            rounds_per_game=1,
+            random_setup_func=None,
+            random_setup_kwargs=None,
+            role_assignator_func=None,
+            role_assignator_func_kwargs=None,
+            finalization_visibility=False,
+            other_values_visibility=False,
+            random_seed=None
+        ):
+            """Initialize the Deal or No Deal environment.
+            
+            Args:
+                agents: List of agent IDs participating in the game
+                mode: Game mode ("coop" or "comp")
+                max_messages: Maximum number of messages per agent per round
+                min_messages: Minimum number of messages per agent per round
+                max_chars_per_message: Maximum characters per message
+                rounds_per_game: Number of negotiation rounds to play
+                random_setup_func: Function to generate item quantities and values
+                random_setup_kwargs: Arguments for the random setup function
+                role_assignator_func: Function to assign roles to agents
+                role_assignator_func_kwargs: Arguments for the role assignator
+                finalization_visibility: Whether agents can see each other's finalizations
+                other_values_visibility: Whether agents can see each other's values
+                random_seed: Seed for reproducibility
+            """
+            # ...
+            
+        def reset(self):
+            """Reset the environment to an initial state and return the initial observation.
+            
+            Returns:
+                observation (dict): A dictionary where keys are agent identifiers and values are observations.
+            """
+            # ...
+            
+        def step(self, actions):
+            """Take a step in the environment using the provided actions.
+
+            Args:
+                actions (dict): A dictionary where keys are agent identifiers and values are actions.
+                    Actions can be messages or finalization proposals.
+
+            Returns:
+                observations (dict): A dictionary where keys are agent identifiers and values are observations.
+                done (bool): Whether the episode has ended.
+                info (dict): Additional information about the environment.
+            """
+            # ...
+            
+        def get_state(self):
+            """Retrieve the current state of the game.
+            
+            Returns:
+                state (dict): The current state of the game, including items, quantities, values, etc.
+            """
+            # ...
+
+Key Implementation Details
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The ``DondEnv`` class implements several key features:
+
+1. **Multi-Agent Support**: The environment tracks two agents and manages their alternating messages.
+
+2. **Turn-Based Dialogue**: The environment enforces turn structure and limits on message count.
+
+3. **Finalization Processing**: The environment validates and processes finalization proposals.
+
+4. **Random Setup**: The environment supports multiple methods of generating negotiation scenarios.
+
+5. **Round Management**: The environment can handle multiple rounds with different setups.
+
+Observation Structure
+~~~~~~~~~~~~~~~~~~~~
+
+Each agent receives an observation (state) dictionary with rich information about the game:
+
+.. code-block:: python
+
+    {
+        "mode": str,                 # Game mode ("coop" or "comp")
+        "role_values": dict,         # Value mappings for each role
+        "role_props": dict,          # Properties for each role
+        "agent_to_role": dict,       # Mapping from agent IDs to roles
+        "is_new_round": bool,        # Whether this is the start of a new round
+        "is_new_game": bool,         # Whether this is the start of a new game
+        "game_over": bool,           # Whether the game is over
+        "items": list,               # List of item names
+        "quantities": dict,          # Quantities of each item
+        "has_finalized": bool,       # Whether finalization has been proposed
+        "last_message": dict,        # The last message sent
+        "messages_remaining": dict,  # Number of messages each agent can still send
+        # And various history tracking fields
+    }
+
+Action Structure
+~~~~~~~~~~~~~~~
+
+Actions can be:
+
+1. **Text Messages**: Free-form text for negotiation.
+2. **Finalization Proposals**: Structured data specifying the exact allocation of items.
+
+Example finalization format:
+
+.. code-block:: python
+
+    {
+        "type": "finalize",
+        "allocation": {
+            "agent1": {"book": 3, "hat": 0, "ball": 6},
+            "agent2": {"book": 1, "hat": 2, "ball": 0}
+        }
+    }
+
+Value Setup Functions
+--------------------
+
+The DoND environment provides several functions for setting up item values:
+
+.. code-block:: python
+
+    def dond_random_setup(items, min_quant, max_quant, min_val, max_val, random_seed=None):
+        """
+        Generates items, even-numbered quantities and distinct random values for each category for both agents.
+        
+        Args:
+            items (list): List of items.
+            min_quant (int): Minimum quantity per item.
+            max_quant (int): Maximum quantity per item.
+            min_val (int): Minimum value per item.
+            max_val (int): Maximum value per item.
+            random_seed (int, optional): Seed for random generation.
+        
+        Returns:
+            tuple: (items, quantities, (val_starting_negotiator, val_responding_negotiator))
+        """
+        # ...
+    
+    def independent_random_vals(items, min_quant, max_quant, min_val, max_val, random_seed=None):
+        """
+        Generates random quantities and independent random values for both agents.
+        
+        Args:
+            Similar to dond_random_setup
+        
+        Returns:
+            tuple: (items, quantities, (val_starting_negotiator, val_responding_negotiator))
+        """
+        # ...
+    
+    def bicameral_vals_assignator(items, min_quant, max_quant, low_val_mean, low_val_std, high_val_mean, high_val_std, random_seed=None):
+        """
+        Generates values with a bicameral distribution - each agent values half the items highly.
+        
+        Args:
+            items (list): List of items.
+            min_quant, max_quant: Range for quantities
+            low_val_mean, low_val_std: Mean and standard deviation for the "low value" distribution
+            high_val_mean, high_val_std: Mean and standard deviation for the "high value" distribution
+            random_seed: Seed for reproducibility
+        
+        Returns:
+            tuple: (items, quantities, (val_starting_negotiator, val_responding_negotiator))
+        """
+        # ...
+
+Running DoND Games
+----------------------
+
+To run Deal or No Deal games with LLM agents, you can use the following structure:
+
+.. code-block:: python
+
+    from mllm.environments.dond.dond_game import DondEnv
+    from mllm.environments.dond.dond_agent import DondAgent
+    from src.run_matches import run_batched_matches
+
+    # Create environment
+    env = DondEnv(
+        agents=["agent1", "agent2"],
+        mode="coop",
+        max_messages=10,
+        rounds_per_game=1,
+        random_setup_func="dond_random_setup",
+        random_setup_kwargs={
+            "items": ["book", "hat", "ball"],
+            "min_quant": 2,
+            "max_quant": 8,
+            "min_val": 1,
+            "max_val": 10
+        },
+        finalization_visibility=False
+    )
+    
+    # Create agent handlers (implementation details would vary)
+    agent_handlers = {
+        "agent1": DondAgent(agent_id="agent1"),
+        "agent2": DondAgent(agent_id="agent2")
+    }
+
+    # Define policy mapping
+    policy_mapping = {
+        "llm_policy": my_llm_policy_function
+    }
+
+    # Run the game
+    game_results = run_batched_matches(
+        envs=[env],
+        agent_handlers_per_env=[agent_handlers],
+        policy_mapping=policy_mapping,
+        max_parallel_matches=1
+    )
+
+Limitations and Considerations
+-----------------------------
+
+1. **Negotiation Complexity**: The open-ended nature of negotiations can be challenging for some LLM agents.
+
+2. **Parsing Challenges**: Extracting structured finalization proposals from free-form text requires robust parsing.
+
+3. **Optimization Opportunities**: Different agents may employ different negotiation strategies to optimize outcomes.
+
+4. **Fairness Evaluation**: The environment allows research into questions of fair division and Pareto optimality.
+
+5. **Strategic Deception**: Agents might strategically misrepresent their true values, adding complexity to negotiations.
+
+Advanced Usage
+------------
+
+For advanced usage, you can:
+
+1. **Custom Value Functions**: Create more complex distributions of item values for specific research questions.
+
+2. **Novel Negotiation Scenarios**: Design item sets and values to test specific negotiation skills.
+
+3. **Curriculum Learning**: Create progressively more difficult negotiation scenarios.
+
+4. **Communication Analysis**: Analyze the language and strategies used in successful negotiations.
+
+5. **Multi-Round Dynamics**: Study how agents adapt their strategies over multiple rounds.

src_code_for_reproducibility/docs/source/environments/ipd.rst

@@ -0,0 +1,411 @@
+=================
+Iterated Prisoner's Dilemma
+=================
+
+The Iterated Prisoner's Dilemma environment provides a classic game theory setting for studying cooperation 
+and competition between agents. This document describes the API for interacting with the IPD environment
+and its associated agent handler.
+
+Overview
+--------
+
+The Prisoner's Dilemma is a fundamental problem in game theory that demonstrates why two rational individuals might not 
+cooperate, even when it appears in their best interest to do so. In the iterated version, the same two players 
+repeatedly face the same dilemma, allowing for the development of trust or retaliation based on previous interactions.
+
+Our implementation follows the Multi-Agent Negotiation Environment standard, allowing it to be used with 
+LLM agents through a text-based interface.
+
+Game Rules
+----------
+
+### Basic Premise
+
+The scenario behind the Prisoner's Dilemma is as follows:
+
+Two criminals are arrested and imprisoned. Each prisoner is in solitary confinement with no means of communicating with 
+the other. The prosecutors lack sufficient evidence to convict the pair on the principal charge, but they have enough 
+to convict both on a lesser charge. Simultaneously, the prosecutors offer each prisoner a bargain:
+
+- If both prisoners betray each other, each serves 2 years in prison (the "punishment" payoff)
+- If one betrays the other while the other remains silent, the betrayer goes free (the "temptation" payoff) while the 
+  silent accomplice serves 3 years (the "sucker" payoff)
+- If both remain silent, each serves only 1 year in prison (the "reward" payoff)
+
+### Game Mechanics
+
+In our implementation, the choices are simplified to:
+- **C**: Cooperate (remain silent)
+- **D**: Defect (betray the other prisoner)
+
+Each round, both players simultaneously choose either C or D, and receive points based on the combination of their choices:
+
+- Both choose C: Both receive the "reward" payoff (3 points by default)
+- Both choose D: Both receive the "punishment" payoff (1 point by default)
+- One chooses C, one chooses D: The defector receives the "temptation" payoff (5 points by default), while the cooperator 
+  receives the "sucker" payoff (0 points by default)
+
+### Example: Single Round
+
+Let's see how a single round plays out:
+
+1. Alice and Bob simultaneously make their choices
+2. If Alice chooses C and Bob chooses C:
+   - Alice receives 3 points
+   - Bob receives 3 points
+3. If Alice chooses C and Bob chooses D:
+   - Alice receives 0 points
+   - Bob receives 5 points
+4. If Alice chooses D and Bob chooses C:
+   - Alice receives 5 points
+   - Bob receives 0 points
+5. If Alice chooses D and Bob chooses D:
+   - Alice receives 1 point
+   - Bob receives 1 point
+
+### Iterated Game Structure
+
+The iterated version repeats this basic game for a fixed number of rounds. The key features are:
+
+1. Players know the total number of rounds in advance
+2. After each round, players learn what choice the other player made
+3. Players maintain a cumulative score across all rounds
+4. Players can adjust their strategy based on the history of previous interactions
+
+### Game Variations
+
+The IPD environment supports several variations through configuration parameters:
+
+#### Different Payoff Matrices
+
+The standard payoff values can be modified to create different incentive structures:
+- **Traditional PD**: reward=3, punishment=1, temptation=5, sucker=0
+- **Weak Temptation**: reward=3, punishment=1, temptation=4, sucker=0 (reduces the incentive to defect)
+- **Harsh Punishment**: reward=3, punishment=0, temptation=5, sucker=0 (increases the cost of mutual defection)
+- **Generous**: reward=4, punishment=2, temptation=5, sucker=1 (cushions the blow of being betrayed)
+
+#### Game Length Variations
+
+The number of rounds can significantly impact strategy:
+- **Short Games** (5-10 rounds): Incentivizes more defection, especially near the end
+- **Medium Games** (20-50 rounds): Allows for the development of tit-for-tat and forgiveness strategies
+- **Long Games** (100+ rounds): Favors steady cooperation with occasional "probing" defections
+
+### Common Strategies
+
+While not enforced by the environment, several well-known strategies can emerge:
+- **Always Cooperate**: Always choose C
+- **Always Defect**: Always choose D
+- **Tit for Tat**: Start with C, then copy what the opponent did in the previous round
+- **Forgiving Tit for Tat**: Like Tit for Tat, but occasionally cooperate even after being defected against
+- **Grudger**: Cooperate until the opponent defects once, then always defect
+- **Random**: Choose randomly between C and D
+
+IPDEnv
+------
+
+The ``IPDEnv`` class provides an interface to the Iterated Prisoner's Dilemma environment that follows the 
+Multi-Agent Negotiation Environment standard.
+
+.. code-block:: python
+
+    class IPDEnv:
+        """
+        Iterated Prisoner's Dilemma environment following the MarlEnvironment standard.
+        
+        In each round of the game, two agents simultaneously choose to either cooperate (C) or defect (D).
+        The payoffs are as follows:
+        - If both cooperate: Both receive the "reward" (usually 3 points)
+        - If both defect: Both receive the "punishment" (usually 1 point)
+        - If one cooperates and one defects: The defector receives the "temptation" (usually 5 points)
+          and the cooperator receives the "sucker" payoff (usually 0 points)
+        
+        The game is played for a specified number of rounds.
+        """
+        
+        def __init__(
+            self,
+            rounds_per_game: int = 10,
+            reward: float = 3.0,           # Both cooperate
+            punishment: float = 1.0,       # Both defect
+            temptation: float = 5.0,       # Defector's reward when other cooperates
+            sucker: float = 0.0,           # Cooperator's reward when other defects
+            random_seed: Optional[int] = None,
+        ):
+            """
+            Initialize the Iterated Prisoner's Dilemma environment.
+            
+            Args:
+                rounds_per_game: Number of rounds to play
+                reward: Payoff when both agents cooperate
+                punishment: Payoff when both agents defect
+                temptation: Payoff for defecting when other agent cooperates
+                sucker: Payoff for cooperating when other agent defects
+                seed: Random seed for reproducibility
+            """
+            # ...
+            
+        def reset(self) -> Dict[str, Dict[str, Any]]:
+            """
+            Reset the environment to an initial state and return the initial observation.
+            
+            Returns:
+                observation (dict): A dictionary where keys are agent identifiers and values are observations.
+            """
+            # ...
+        
+        def step(self, actions: Dict[str, str]) -> Tuple[Dict[str, Dict[str, Any]], bool, Dict[str, Any]]:
+            """
+            Take a step in the environment using the provided actions.
+            
+            Args:
+                actions (dict): A dictionary where keys are agent identifiers and values are actions ('C' or 'D').
+            
+            Returns:
+                observations (dict): A dictionary where keys are agent identifiers and values are observations.
+                done (bool): Whether the episode has ended.
+                info (dict): Additional information about the environment.
+            """
+            # ...
+
+Key Implementation Details
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The ``IPDEnv`` class implements several key features:
+
+1. **Two-Agent Support**: The environment tracks two agents ("alice" and "bob") and manages their interactions.
+
+2. **Round-Based Play**: The environment enforces turn structure and tracks game history.
+
+3. **Payoff Matrix**: The environment calculates rewards based on the standard prisoner's dilemma payoff matrix.
+
+4. **Observation Generation**: The environment generates detailed observations for each agent, including action history and rewards.
+
+5. **Game Termination**: The environment tracks game termination after the specified number of rounds.
+
+Observation Structure
+~~~~~~~~~~~~~~~~~~~~
+
+Each agent receives an observation dictionary with the following structure:
+
+.. code-block:: python
+
+    {
+        "current_round": int,                # Current round number (0-indexed)
+        "rounds_per_game": int,              # Total number of rounds in the game
+        "history": List[Dict],               # Complete game history so far
+        "last_round_actions": Dict[str, str], # Actions from the previous round (if any)
+        "last_round_reward": float,          # Reward received in the previous round (if any)
+        "total_reward": float,               # Cumulative reward so far
+        "payoff_matrix": Dict[str, float],   # The game's payoff matrix values
+    }
+
+Action Structure
+~~~~~~~~~~~~~~~
+
+Actions are simple strings:
+
+1. ``"C"`` for Cooperate
+2. ``"D"`` for Defect
+
+IPDAgent
+--------------
+
+The ``IPDAgent`` class implements the agent handler interface for the Iterated Prisoner's Dilemma, processing observations from the environment and generating actions through an LLM.
+
+.. code-block:: python
+
+    class IPDAgent:
+        """
+        Agent handler for Iterated Prisoner's Dilemma, implementing the AgentState interface 
+        for the multi-agent negotiation standard.
+        """
+        
+        def __init__(
+            self,
+            agent_id: str,
+            policy_id: str = "llm_policy",
+            system_prompt: Optional[str] = None,
+            max_errors: int = 3,
+            opponent_id: Optional[str] = None,
+        ):
+            """
+            Initialize the IPD agent handler.
+            
+            Args:
+                agent_id: Identifier for this agent ("alice" or "bob")
+                policy_id: Identifier for the policy this agent uses
+                system_prompt: Optional custom system prompt for the LLM
+                max_errors: Maximum number of parsing errors before defaulting to cooperate
+                opponent_id: Optional identifier of the opponent (inferred if not provided)
+            """
+            # ...
+            
+        def step(self, observation_from_env: Dict[str, Any], policy_output: str = None) -> Tuple[str, Dict[str, Any], str, bool, Dict[str, Any]]:
+            """
+            Update the agent state based on the observation and process the policy output.
+            
+            Args:
+                observation_from_env: The observation from the environment
+                policy_output: The output from the policy (LLM response)
+                
+            Returns:
+                policy_id: The policy identifier
+                policy_input: The input to the policy
+                action: The action to be sent to the environment
+                done: Whether the action is ready to be sent to the environment
+                info: Additional information about the agent
+            """
+            # ...
+
+Key Implementation Details
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The ``IPDAgent`` class implements several key features:
+
+1. **LLM Interaction**: The agent generates prompts for an LLM and processes the LLM's responses.
+
+2. **Action Extraction**: The agent parses the LLM's output to extract valid actions (C or D).
+
+3. **Error Handling**: The agent provides helpful error messages when parsing fails and defaults to cooperation after multiple failures.
+
+4. **History Tracking**: The agent maintains and provides the complete game history in its prompts.
+
+5. **Strategy Explanation**: The agent can extract and log the reasoning behind an LLM's decisions.
+
+Prompt Structure
+~~~~~~~~~~~~~~~
+
+The agent generates prompts that include:
+
+1. **System Prompt**: Instructions and context for the LLM, explaining its role and the rules of the Prisoner's Dilemma.
+
+2. **Game State Description**: A text description of the current game state, including:
+   - Current round number
+   - History of previous rounds (if any)
+   - Cumulative score
+
+3. **Action Request**: Instructions on how to format the response, requiring an explicit action tag.
+
+Example system prompt:
+
+.. code-block:: text
+
+    You are playing as Alice in an Iterated Prisoner's Dilemma game against Bob.
+    In each round, you must choose to either Cooperate (C) or Defect (D).
+    
+    The payoffs are:
+    - If both players Cooperate: You each get 3 points
+    - If both players Defect: You each get 1 point
+    - If you Cooperate and Bob Defects: You get 0 points, Bob gets 5 points
+    - If you Defect and Bob Cooperates: You get 5 points, Bob gets 0 points
+    
+    Your goal is to maximize your total points across all rounds.
+    The game will last for exactly 10 rounds, and both players know this.
+
+Example game state prompt:
+
+.. code-block:: text
+
+    Current round: 3/10
+    
+    History:
+    Round 1: You chose C, Bob chose C. You earned 3 points.
+    Round 2: You chose C, Bob chose D. You earned 0 points.
+    
+    Your total score so far: 3 points
+    
+    What is your choice for round 3?
+    Please respond with <action>C</action> to cooperate or <action>D</action> to defect,
+    and explain your reasoning.
+
+Running IPD Games
+----------------------
+
+To run Iterated Prisoner's Dilemma games with LLM agents, you can use the following code structure:
+
+.. code-block:: python
+
+    from mllm.environments.ipd.ipd_game import IPDEnv
+    from mllm.environments.ipd.ipd_agent import IPDAgent
+    from mllm.run_matches import run_batched_matches
+
+    # Create environment
+    env = IPDEnv(
+        rounds_per_game=10,
+        reward=3.0,
+        punishment=1.0,
+        temptation=5.0,
+        sucker=0.0
+    )
+    
+    # Create agent handlers
+    agent_handlers = {
+        "alice": IPDAgent(agent_id="alice"),
+        "bob": IPDAgent(agent_id="bob")
+    }
+
+    # Define policy mapping
+    policy_mapping = {
+        "llm_policy": my_llm_policy_function
+    }
+
+    # Run the game
+    game_results = run_batched_matches(
+        envs=[env],
+        agent_handlers_per_env=[agent_handlers],
+        policy_mapping=policy_mapping,
+        max_parallel_matches=1
+    )
+
+    # Process results
+    for result in game_results:
+        print(f"Game finished. Scores: {result['total_rewards']}")
+
+Statistics and Analysis
+----------------------
+
+The IPD environment includes utility functions for analyzing game outcomes:
+
+1. **Cooperation Rates**: Percentage of rounds where each agent cooperated.
+2. **Mutual Cooperation/Defection**: Percentage of rounds where both agents made the same choice.
+3. **Score Distribution**: Analysis of how points were accumulated over the game.
+
+These statistics can be calculated using the ``gather_ipd_statistics`` function:
+
+.. code-block:: python
+
+    from mllm.environments.ipd.ipd_statistics_funcs import gather_ipd_statistics
+    
+    stats = gather_ipd_statistics(match_info, env_info)
+    print(f"Cooperation rates: {stats['cooperation_rate']}")
+    print(f"Mutual cooperation rate: {stats['mutual_cooperation_rate']}")
+    print(f"Mutual defection rate: {stats['mutual_defection_rate']}")
+
+Limitations and Considerations
+-----------------------------
+
+1. **Determinism**: The environment is deterministic, with randomness only in initialization if a seed is provided.
+
+2. **Limited Player Count**: The IPD environment only supports exactly two players.
+
+3. **Perfect Information**: Both players have perfect information about the game history.
+
+4. **Simultaneous Actions**: Both players act simultaneously, which requires adaptations for some LLM interfaces.
+
+5. **Fixed Game Length**: The total number of rounds is fixed and known to both players from the start.
+
+Advanced Usage
+------------
+
+For advanced usage, you can customize:
+
+1. **Payoff Matrix**: Modify reward values to create different incentive structures.
+
+2. **System Prompts**: Customize the LLM's understanding of the game and potential strategies.
+
+3. **Error Handling**: Adjust how the agent responds to invalid LLM outputs.
+
+4. **Analysis**: Create custom statistics gathering for specific research questions.
+
+5. **Integration**: Connect the IPD environment to other negotiation frameworks or tournament systems.

src_code_for_reproducibility/docs/source/src.environments.dond.dond_game.rst

@@ -0,0 +1,7 @@
+src.environments.dond.dond\_game module
+=======================================
+
+.. automodule:: src.environments.dond.dond_game
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.environments.dond.dond_log_funcs.rst

@@ -0,0 +1,7 @@
+src.environments.dond.dond\_log\_funcs module
+=============================================
+
+.. automodule:: src.environments.dond.dond_log_funcs
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.environments.dond.dond_player.rst

@@ -0,0 +1,7 @@
+src.environments.dond.dond\_agent module
+=========================================
+
+.. automodule:: src.environments.dond.dond_agent
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.environments.env_imports.rst

@@ -0,0 +1,7 @@
+src.environments.env\_imports module
+====================================
+
+.. automodule:: src.environments.env_imports
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.environments.ipd.ipd_log_funcs.rst

@@ -0,0 +1,7 @@
+src.environments.ipd.ipd\_log\_funcs module
+===========================================
+
+.. automodule:: src.environments.ipd.ipd_log_funcs
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.environments.ipd.rst

@@ -0,0 +1,19 @@
+src.environments.ipd package
+============================
+
+.. automodule:: src.environments.ipd
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Submodules
+----------
+
+.. toctree::
+   :maxdepth: 4
+
+   src.environments.ipd.ipd_agent
+   src.environments.ipd.ipd_game
+   src.environments.ipd.ipd_log_funcs
+   src.environments.ipd.ipd_statistics_funcs
+   src.environments.ipd.ipd_training_data_funcs

src_code_for_reproducibility/docs/source/src.environments.rst

@@ -0,0 +1,25 @@
+src.environments package
+========================
+
+.. automodule:: src.environments
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Subpackages
+-----------
+
+.. toctree::
+   :maxdepth: 4
+
+   src.environments.dond
+   src.environments.ipd
+
+Submodules
+----------
+
+.. toctree::
+   :maxdepth: 4
+
+   src.environments.env_imports
+   src.environments.environment_imports

src_code_for_reproducibility/docs/source/src.experiments.dond_run_train.rst

@@ -0,0 +1,7 @@
+src.experiments.dond\_run\_train module
+=======================================
+
+.. automodule:: src.experiments.dond_run_train
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.generation.run_games.rst

@@ -0,0 +1,7 @@
+src.generation.run\_games module
+================================
+
+.. automodule:: src.generation.run_games
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.models.hf_agent.rst

@@ -0,0 +1,7 @@
+src.models.hf\_agent module
+===========================
+
+.. automodule:: src.models.hf_agent
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.models.vllm_worker_wrap.rst

@@ -0,0 +1,7 @@
+src.models.vllm\_worker\_wrap module
+====================================
+
+.. automodule:: src.models.vllm_worker_wrap
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.run.rst

@@ -0,0 +1,7 @@
+src.run module
+==============
+
+.. automodule:: src.run
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.training.reinforce_training.rst

@@ -0,0 +1,7 @@
+src.training.reinforce\_training module
+=======================================
+
+.. automodule:: src.training.reinforce_training
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.training.rl_convs_processing.rst

@@ -0,0 +1,7 @@
+src.training.rl\_convs\_processing module
+=========================================
+
+.. automodule:: src.training.rl_convs_processing
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.utils.inherit_args.rst

@@ -0,0 +1,7 @@
+src.utils.inherit\_args module
+==============================
+
+.. automodule:: src.utils.inherit_args
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.utils.log_gpu_usage.rst

@@ -0,0 +1,7 @@
+src.utils.log\_gpu\_usage module
+================================
+
+.. automodule:: src.utils.log_gpu_usage
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.utils.log_statistics.rst

@@ -0,0 +1,7 @@
+src.utils.log\_statistics module
+================================
+
+.. automodule:: src.utils.log_statistics
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.utils.parallel_shuffle.rst

@@ -0,0 +1,7 @@
+src.utils.parallel\_shuffle module
+==================================
+
+.. automodule:: src.utils.parallel_shuffle
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.utils.quick_stats.rst

@@ -0,0 +1,7 @@
+src.utils.quick\_stats module
+=============================
+
+.. automodule:: src.utils.quick_stats
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/markov_games/alternative_actions_runner.py

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

src_code_for_reproducibility/markov_games/diplomacy/diplomacy_env.py

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

src_code_for_reproducibility/markov_games/gather_and_export_utils.py

@@ -0,0 +1,951 @@
+from __future__ import annotations
+
+import csv
+import os
+import pickle
+import re
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
+
+from mllm.markov_games.rollout_tree import *
+
+try:
+    # Re-export moved helpers for backward compatibility
+    from basic_render import (
+        find_iteration_folders,
+        gather_rollout_trees,
+        get_rollout_trees,
+    )
+except Exception:
+    pass
+
+# --------------------------------------------------------------------------------------
+# Fetch external rollout trees
+# --------------------------------------------------------------------------------------
+
+
+def find_iteration_folders(global_folder):
+    """Find all iteration_* folders within the global folder structure."""
+    global_path = Path(global_folder)
+
+    # Look for iteration_* folders in all subdirectories
+    iteration_folders = []
+
+    # Search in the global folder itself
+    for item in global_path.glob("iteration_*"):
+        if item.is_dir():
+            iteration_folders.append(item)
+
+    # Search in seed_* subdirectories
+    for seed_dir in global_path.glob("seed_*/"):
+        if seed_dir.is_dir():
+            for item in seed_dir.glob("iteration_*"):
+                if item.is_dir():
+                    iteration_folders.append(item)
+
+    return sorted(iteration_folders)
+
+
+def gather_rollout_trees(iteration_folder):
+    """Gather all rollout trees from the iteration folder (.pkl only)."""
+    rollout_trees = []
+    iteration_path = Path(iteration_folder)
+    for item in iteration_path.glob("**/*.rt.pkl"):
+        with open(item, "rb") as f:
+            data = pickle.load(f)
+        # Validate dicts back into Pydantic model for downstream use
+        rollout_tree = RolloutTreeRootNode.model_validate(data)
+        rollout_trees.append(rollout_tree)
+    return rollout_trees
+
+
+def get_rollout_trees(global_folder) -> list[list[RolloutTreeRootNode]]:
+    """Get all rollout trees from the global folder."""
+    iteration_folders = find_iteration_folders(global_folder)
+    rollout_trees = []
+    for iteration_folder in iteration_folders:
+        rollout_trees.append(gather_rollout_trees(iteration_folder))
+    return rollout_trees
+
+
+# --------------------------------------------------------------------------------------
+# Gather data from rollout tree methods
+# --------------------------------------------------------------------------------------
+
+
+def load_rollout_tree(path: Path) -> RolloutTreeRootNode:
+    """Load a rollout tree from a PKL file containing a dict."""
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+    return RolloutTreeRootNode.model_validate(data)
+
+
+@dataclass
+class RolloutNodeList:
+    id: str
+    nodes: List[RolloutTreeNode]
+
+
+def get_rollout_tree_paths(
+    root: RolloutTreeRootNode, mgid: Optional[str] = None
+) -> Tuple[RolloutNodeList, List[RolloutNodeList]]:
+    """
+    Returns:
+        main_path: The main path from the root to the end of the tree.
+        branch_paths: A list of all branch paths from the root to the end of the tree.
+        Each branch path contains a list of nodes that are part of the branch, including the nodes from the main path before the branch was taken.
+    """
+    branch_paths = []
+
+    def collect_path_nodes(current) -> List[RolloutTreeNode]:
+        """Recursively collect all nodes in a path starting from current node."""
+        if current is None:
+            return []
+
+        if isinstance(current, RolloutTreeNode):
+            return [current] + collect_path_nodes(current.child)
+
+        elif isinstance(current, RolloutTreeBranchNode):
+            # For branch nodes, we only follow the main_child for path collection
+            if current.main_child:
+                return [current.main_child] + collect_path_nodes(
+                    current.main_child.child
+                )
+            else:
+                return []
+
+    def traverse_for_branches(
+        current,
+        main_path_prefix: List[RolloutTreeNode],
+        path_id: str,
+        current_time_step: Optional[int] = 0,
+    ):
+        """Traverse tree to collect all branch paths."""
+        if current is None:
+            return
+
+        if isinstance(current, RolloutTreeNode):
+            # Continue traversing with this node added to the main path prefix
+            new_prefix = main_path_prefix + [current]
+            traverse_for_branches(current.child, new_prefix, path_id, current.time_step)
+
+        elif isinstance(current, RolloutTreeBranchNode):
+            # Collect all branch paths
+            if current.branches:
+                for agent_id, branch_node_list in current.branches.items():
+                    if branch_node_list:
+                        # Start with the main path prefix, then recursively collect all nodes in this branch
+                        branch_path_nodes = main_path_prefix.copy()
+                        for branch_node in branch_node_list:
+                            branch_path_nodes.extend(collect_path_nodes(branch_node))
+
+                        # Create proper branch path ID with mgid, agent_id, and time_step
+                        mgid_str = mgid or str(root.id)
+                        branch_path_id = f"mgid:{mgid_str}_type:branch_agent:{agent_id}_time_step:{current_time_step}"
+                        branch_paths.append(
+                            RolloutNodeList(id=branch_path_id, nodes=branch_path_nodes)
+                        )
+
+            # Process the main child and add to prefix
+            new_prefix = main_path_prefix
+            if current.main_child:
+                new_prefix = main_path_prefix + [current.main_child]
+
+            # Continue traversing the main path
+            if current.main_child:
+                traverse_for_branches(
+                    current.main_child.child,
+                    new_prefix,
+                    path_id,
+                    current.main_child.time_step,
+                )
+
+    # Collect the main path nodes
+    main_path_nodes = collect_path_nodes(root.child)
+
+    # Traverse to collect all branch paths
+    traverse_for_branches(root.child, [], "")
+
+    # Create the main path with proper mgid format
+    mgid_str = mgid or str(root.id)
+    main_path = RolloutNodeList(id=f"mgid:{mgid_str}_type:main", nodes=main_path_nodes)
+
+    return main_path, branch_paths
+
+
+class ChatTurnLog(BaseModel):
+    time_step: int
+    agent_id: str
+    role: str
+    content: str
+    reasoning: Optional[str] = None
+    is_state_end: bool
+    reward: float
+
+
+def gather_agent_chat_turns_for_path(
+    agent_id: str, path: RolloutNodeList
+) -> List[ChatTurnLog]:
+    """Iterate through all chat turns for a specific agent in a path sorted by time step."""
+    turns = []
+    for node in path.nodes:
+        action_log = node.step_log.action_logs.get(agent_id, [])
+        if action_log:
+            for chat_turn in action_log.chat_turns or []:
+                turns.append(
+                    ChatTurnLog(
+                        time_step=node.time_step,
+                        agent_id=agent_id,
+                        role=chat_turn.role,
+                        content=chat_turn.content,
+                        reasoning=getattr(chat_turn, "reasoning_content", None),
+                        is_state_end=chat_turn.is_state_end,
+                        reward=node.step_log.simulation_step_log.rewards.get(
+                            agent_id, 0
+                        ),
+                    )
+                )
+    return turns
+
+
+def gather_all_chat_turns_for_path(path: RolloutNodeList) -> List[ChatTurnLog]:
+    """Iterate through all chat turns for all agents in a path sorted by time step."""
+    turns = []
+
+    # Collect turns from all agents, but interleave them per timestep by (user, assistant) pairs
+    for node in path.nodes:
+        # Build (user[, assistant]) pairs for each agent at this timestep
+        agent_ids = sorted(list(node.step_log.action_logs.keys()))
+        per_agent_pairs: Dict[str, List[List[ChatTurnLog]]] = {}
+
+        for agent_id in agent_ids:
+            action_log = node.step_log.action_logs.get(agent_id)
+            pairs: List[List[ChatTurnLog]] = []
+            current_pair: List[ChatTurnLog] = []
+
+            if action_log and action_log.chat_turns:
+                for chat_turn in action_log.chat_turns:
+                    turn_log = ChatTurnLog(
+                        time_step=node.time_step,
+                        agent_id=agent_id,
+                        role=chat_turn.role,
+                        content=chat_turn.content,
+                        reasoning=getattr(chat_turn, "reasoning_content", None),
+                        is_state_end=chat_turn.is_state_end,
+                        reward=node.step_log.simulation_step_log.rewards.get(
+                            agent_id, 0
+                        ),
+                    )
+
+                    if chat_turn.role == "user":
+                        # If a previous pair is open, close it and start a new one
+                        if current_pair:
+                            pairs.append(current_pair)
+                            current_pair = []
+                        current_pair = [turn_log]
+                    else:
+                        # assistant: attach to an open user message if present; otherwise stand alone
+                        if (
+                            current_pair
+                            and len(current_pair) == 1
+                            and current_pair[0].role == "user"
+                        ):
+                            current_pair.append(turn_log)
+                            pairs.append(current_pair)
+                            current_pair = []
+                        else:
+                            # No preceding user or already paired; treat as its own unit
+                            pairs.append([turn_log])
+
+                if current_pair:
+                    # Unpaired trailing user message
+                    pairs.append(current_pair)
+
+            per_agent_pairs[agent_id] = pairs
+
+        # Interleave pairs across agents: A1, B1, A2, B2, ...
+        index = 0
+        while True:
+            added_any = False
+            for agent_id in agent_ids:
+                agent_pairs = per_agent_pairs.get(agent_id, [])
+                if index < len(agent_pairs):
+                    for tl in agent_pairs[index]:
+                        turns.append(tl)
+                    added_any = True
+            if not added_any:
+                break
+            index += 1
+
+    return turns
+
+
+def chat_turns_to_dict(chat_turns: Iterator[ChatTurnLog]) -> Iterator[Dict[str, Any]]:
+    """Render all chat turns for a path as structured data for JSON."""
+    for chat_turn in chat_turns:
+        yield chat_turn.model_dump()
+
+
+def get_all_agents(root: RolloutTreeRootNode) -> List[str]:
+    """list of all agent IDs that appear in the tree."""
+    if root.child is None:
+        return []
+
+    # Get the first node to extract all agent IDs
+    first_node = root.child
+    if isinstance(first_node, RolloutTreeBranchNode):
+        first_node = first_node.main_child
+
+    if first_node is None:
+        return []
+
+    # All agents should be present in the first node
+    agents = set(first_node.step_log.action_logs.keys())
+    agents.update(first_node.step_log.simulation_step_log.rewards.keys())
+
+    return sorted(list(agents))
+
+
+def gather_agent_main_rewards(agent_id: str, path: RolloutNodeList) -> List[float]:
+    """Gather main rewards for a specific agent in a path."""
+    rewards = []
+    for node in path.nodes:
+        reward = node.step_log.simulation_step_log.rewards[agent_id]
+        rewards.append(reward)
+    return rewards
+
+
+def gather_all_rewards(path: RolloutNodeList) -> List[Dict[AgentId, float]]:
+    """Gather main rewards from main trajectory in a path."""
+    rewards = []
+    for node in path.nodes:
+        rewards.append(node.step_log.simulation_step_log.rewards.copy())
+    return rewards
+
+
+def gather_simulation_stats(
+    path: RolloutNodeList,
+    filter: Callable[[SimulationStepLog], bool],
+    stat_func: Callable[[SimulationStepLog], Any],
+) -> List[Any]:
+    """Gather stats from main trajectory in a path."""
+    stats = []
+    for node in path.nodes:
+        sl = node.step_log.simulation_step_log
+        if filter(sl):
+            stats.append(stat_func(sl))
+    return stats
+
+
+def gather_simulation_infos(path: RolloutNodeList) -> List[Dict[str, Any]]:
+    """Gather simulation information from main trajectory in a path."""
+    infos = []
+    for node in path.nodes:
+        infos.append(node.step_log.simulation_step_log.info)
+    return infos
+
+
+def export_chat_logs(path: Path, outdir: Path):
+    """Process a rollout tree PKL file and generate a JSONL of chat turns as dicts.
+    Each line contains an object with path_id and chat_turns for a single path.
+    """
+    import json
+
+    root = load_rollout_tree(path)
+    mgid = root.id
+
+    main_path, branch_paths = get_rollout_tree_paths(root)
+    all_paths = [main_path] + branch_paths
+
+    outdir.mkdir(parents=True, exist_ok=True)
+    output_file = outdir / f"mgid:{mgid}_plucked_chats.render.jsonl"
+
+    with open(output_file, "w", encoding="utf-8") as f:
+        for path_obj in all_paths:
+            chat_turns = gather_all_chat_turns_for_path(path_obj)
+            output_obj = {
+                "path_id": str(path_obj.id),
+                "chat_turns": list(chat_turns_to_dict(iter(chat_turns))),
+            }
+            f.write(json.dumps(output_obj, ensure_ascii=False) + "\n")
+
+
+def export_rewards_to_csv(path: Path, outdir: Path, first_file: bool):
+    # Load the rollout tree
+    root = load_rollout_tree(path)
+    mgid = root.id
+
+    # Get all paths
+    main_path, branch_paths = get_rollout_tree_paths(root)
+    outdir.mkdir(parents=True, exist_ok=True)
+    rewards_dict_list = gather_all_rewards(main_path)
+    agent_ids = rewards_dict_list[0].keys()
+    rewards_list = defaultdict(list)
+    for rewards_dict in rewards_dict_list:
+        for agent_id in agent_ids:
+            rewards_list[agent_id].append(rewards_dict[agent_id])
+
+    mgid = root.id
+    group_seed = getattr(root, "crn_id", None)
+
+    for agent_id in agent_ids:
+        output_file = outdir / f"agent:{agent_id}_rewards.render.csv"
+
+        # Build current row: [mgid, group_seed] + rewards
+        formatted_rewards = [f"{round(x, 1):>5}" for x in rewards_list[agent_id]]
+        current_row = [str(mgid), str(group_seed)] + formatted_rewards
+
+        # Read existing rows (if any), skipping header if present
+        existing_rows: List[List[str]] = []
+        if output_file.exists():
+            with open(output_file, "r", newline="") as rf:
+                reader = csv.reader(rf)
+                for row in reader:
+                    if not row or not any(cell.strip() for cell in row):
+                        continue
+                    if (
+                        len(row) >= 2
+                        and row[0].strip().lower() == "mgid"
+                        and row[1].strip().lower() == "group_seed"
+                    ):
+                        # skip header
+                        continue
+                    existing_rows.append(row)
+
+        # Append and sort by (group_seed, mgid)
+        existing_rows.append(current_row)
+
+        def sort_key(r: List[str]):
+            def try_int(val: str):
+                try:
+                    return int(val)
+                except Exception:
+                    return None
+
+            seed_raw = r[1] if len(r) > 1 else ""
+            mgid_raw = r[0] if len(r) > 0 else ""
+            seed_num = try_int(seed_raw)
+            mgid_num = try_int(mgid_raw)
+            # Sort numerically when possible; otherwise fall back to string
+            return (
+                0 if seed_num is not None else 1,
+                seed_num if seed_num is not None else seed_raw,
+                0 if mgid_num is not None else 1,
+                mgid_num if mgid_num is not None else mgid_raw,
+            )
+
+        existing_rows.sort(key=sort_key)
+
+        # Determine max reward length to build header and pad rows
+        max_reward_len = 0
+        for r in existing_rows:
+            if len(r) > 2:
+                max_reward_len = max(max_reward_len, len(r) - 2)
+        max_reward_len = max(max_reward_len, len(current_row) - 2)
+
+        def pad_row(r: List[str]) -> List[str]:
+            needed = (2 + max_reward_len) - len(r)
+            return r + ([""] * needed if needed > 0 else [])
+
+        padded_rows = [pad_row(r) for r in existing_rows]
+
+        # Build header
+        header = ["mgid", "group_seed"] + [f"r_t{t}" for t in range(max_reward_len)]
+
+        # Rewrite the file with header to avoid extra/blank rows
+        with open(output_file, "w", newline="") as wf:
+            writer = csv.writer(wf)
+            writer.writerow(header)
+            writer.writerows(padded_rows)
+
+
+# --------------------------------------------------------------------------------------
+# HTML exports
+# --------------------------------------------------------------------------------------
+
+
+def html_from_chat_turns(chat_turns: List[ChatTurnLog]) -> str:
+    """
+    Render chat turns as a single, wrapping sequence of messages in time order.
+    Keep badge and message bubble styles, include time on every badge and
+    include rewards on assistant badges. Each message is individually
+    hide/show by click; when hidden, only the badge remains and "(...)" is
+    shown inline (not inside a bubble).
+    """
+    import html
+
+    # Prepare ordering: sort by (time_step, original_index) to keep stable order within same step
+    indexed_turns = list(enumerate(chat_turns))
+    indexed_turns.sort(key=lambda t: (t[1].time_step, t[0]))
+
+    # CSS styles (simplified layout; no time-step or agent-column backgrounds)
+    css = """
+    <style>
+        :root {
+            --font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+            --bg: #ffffff;
+            --text: #1c0b00;
+            --muted-text: #2C3E50;
+            --accent-muted: #BDC3C7;
+            --accent-muted-2: #D0D7DE;
+            --panel-bg: #F8FAFC;
+            --reward-color: #a89206;
+            --font-size: 15px;
+            --small-font-size: 13px;
+            --group-label-font-size: 12px;
+            --border-width: 2px;
+            --corner-radius: 6px;
+            --pill-radius-left: 999px 0 0 999px;
+            --pill-radius-right: 0 999px 999px 0;
+            --inset-shadow: 0 1px 0 rgba(0,0,0,0.03) inset;
+        }
+        body {
+            font-family: var(--font-family);
+            margin: 16px;
+            background-color: var(--bg);
+            color: var(--text);
+            font-size: var(--font-size);
+            line-height: 1.6;
+        }
+        .messages-flow {
+            display: block; /* behave like a text container */
+        }
+        .toolbar {
+            display: flex;
+            align-items: center;
+            gap: 8px;
+            margin-bottom: 0;
+            font-size: var(--small-font-size);
+            max-height: 0;
+            overflow: hidden;
+            opacity: 0;
+            pointer-events: none;
+            transition: max-height 0.2s ease, opacity 0.2s ease;
+        }
+        .toolbar-wrap { position: sticky; top: 0; z-index: 10; background: var(--bg); }
+        .toolbar-hotzone { height: 6px; }
+        .toolbar-wrap:hover .toolbar { max-height: 200px; opacity: 1; pointer-events: auto; margin-bottom: 12px; }
+        .toolbar input[type="number"] {
+            width: 72px;
+            padding: 2px 6px;
+            border: 1px solid var(--accent-muted);
+            border-radius: var(--corner-radius);
+            background: var(--bg);
+        }
+        .toolbar button {
+            padding: 4px 8px;
+            border: 1px solid var(--accent-muted);
+            background: var(--panel-bg);
+            border-radius: var(--corner-radius);
+            cursor: pointer;
+        }
+        .chat-turn {
+            display: inline; /* inline like text */
+            background: transparent;
+            position: relative;
+            cursor: default;
+        }
+        /* No agent-specific background distinctions */
+        .turn-content {
+            white-space: normal;
+            color: var(--text);
+            font-size: var(--font-size);
+            display: inline; /* inline flow */
+        }
+        .chat-turn .agent-badge { margin-right: 0; vertical-align: baseline; }
+        .agent-badge {
+            display: inline;
+            position: relative;
+            border: var(--border-width) solid var(--accent-muted); /* slightly thicker */
+            border-radius: var(--pill-radius-left); /* round left and bottom-right */
+            font-size: var(--font-size);
+            color: var(--muted-text);
+            background: transparent;
+            box-shadow: var(--inset-shadow);
+            line-height: 1.35;
+            padding: 3px 10px;
+            border-right: 0;
+            cursor: default;
+        }
+        .agent-badge::after {
+            content: none;
+        }
+        /* removed external separator; emoji is rendered inside message bubble */
+        .agent-name { font-weight: 700; }
+        .emoji-bw { filter: grayscale(100%); opacity: 0.95; font-size: var(--font-size); vertical-align: baseline; margin: 0; position: relative; top: -1px; line-height: 1; display: inline-block; }
+        .ts-badge {
+            position: relative;
+            display: inline;
+            border: var(--border-width) solid var(--accent-muted-2); /* slightly thicker */
+            border-radius: var(--corner-radius); /* not a pill */
+            font-size: var(--font-size);
+            font-weight: 700;
+            color: var(--muted-text);
+            background: #F4F8FB; /* subtle tint */
+            padding: 1px 6px; /* slight padding for visibility */
+            margin-right: 8px; /* small gap from following content */
+            pointer-events: auto; /* allow events so we can ignore them in JS */
+        }
+        /* Hide timestep badges when grouping by 1 */
+        .hide-ts-badges .ts-badge { display: none; }
+        /* Strong hide: completely hide collapsed segments */
+        .strong-hide .segment.collapsed { display: none; }
+        .ts-badge::before {
+            content: "";
+            position: relative;
+            background: var(--accent-muted-2);
+            border-radius: 2px;
+        }
+        .agent-badge { margin-left: 6px;  }
+        /* Segments (reasoning and message) */
+        .segment {
+            display: inline; /* inline bubble behaving like text */
+            font-size: var(--font-size);
+            position: relative;
+            background: var(--bg);
+            vertical-align: baseline;
+            line-height: 1.35;
+            cursor: pointer;
+        }
+        .message-box, .reasoning-box {
+            display: inline; /* inline bubble behaving like text */
+            font-size: var(--font-size);
+            border: var(--border-width) solid var(--accent-muted);
+            border-radius: var(--pill-radius-right); /* message defaults to pill-right */
+            position: relative;
+            background: var(--bg);
+            vertical-align: baseline;
+            line-height: 1.35;
+            padding: 3px 10px;
+            border-left: 0;
+        }
+        /* Reasoning between badge and message: no left or right rounding, seam on both sides */
+        .reasoning-box {
+            border-radius: 0;
+            border-left: 0;
+            border-right: 0;
+        }
+        /* Reasoning text style: slightly smaller and slightly gray */
+        .reasoning-box .seg-text {
+            font-size: var(--small-font-size);
+            color: #6b7280;
+        }
+        .message-box::before { content: none; display: none; margin-right: 0; line-height: 1; }
+        .reasoning-box::before { content: none; display: none; margin-right: 0; line-height: 1; }
+        /* Segment collapsed behavior */
+        .segment .seg-text { display: inline; }
+        .segment.collapsed .seg-text { color: transparent; font-size: 0; display: inline-block; }
+        .segment.collapsed::after { content: "(...)"; color: #7f8c8d; font-style: italic; font-size: var(--font-size); line-height: 1.2; }
+        .segment.collapsed .emoji-bw { opacity: 0.3; }
+        .chat-turn.agent-alice.role-assistant .message-box::before { color: #0eb224; }
+        .chat-turn.agent-bob.role-assistant .message-box::before { color: #ef8323; }
+        .chat-turn.collapsed .message-box::before { display: none; }
+        /* Assistant bubble border colors by common agent names */
+        .chat-turn.agent-alice.role-assistant .message-box { border-color: #0eb224; }
+        .chat-turn.agent-bob.role-assistant .message-box { border-color: #ef8323; }
+        .chat-turn.agent-alice.role-assistant .reasoning-box { border-color: #0eb224; }
+        .chat-turn.agent-bob.role-assistant .reasoning-box { border-color: #ef8323; }
+        /* Tie badge and seam to agent color for a cohesive capsule, assistants only */
+        .chat-turn.agent-alice.role-assistant .agent-badge { border-color: #0eb224; }
+        .chat-turn.agent-alice.role-assistant .agent-badge::after { border-right-color: #0eb224; }
+        .chat-turn.agent-alice.role-assistant .turn-content::before { border-left-color: #0eb224; border-top-color: #0eb224; }
+        .chat-turn.agent-alice.role-assistant .message-box { border-color: #0eb224; }
+
+        .chat-turn.agent-bob.role-assistant .agent-badge { border-color: #ef8323; }
+        .chat-turn.agent-bob.role-assistant .agent-badge::after { border-right-color: #ef8323; }
+        .chat-turn.agent-bob.role-assistant .turn-content::before { border-left-color: #ef8323; border-top-color: #ef8323; }
+        .chat-turn.agent-bob.role-assistant .message-box { border-color: #ef8323; }
+        /* No colored agent-name; keep neutral */
+        .reward { color: var(--reward-color); font-weight: 600; } /* dark gold */
+        .message-placeholder { display: none; color: #7f8c8d; font-style: italic; }
+        /* Group divider - clearer and pretty */
+        .group-divider {
+            display: flex;
+            align-items: center;
+            gap: 8px;
+            width: 100%;
+            margin: 8px 0 2px 0;
+        }
+        .group-divider::before,
+        .group-divider::after {
+            content: "";
+            flex: 1 1 auto;
+            height: 2px;
+            background: linear-gradient(90deg, rgba(224,230,235,0), var(--accent-muted-2) 30%, var(--accent-muted-2) 70%, rgba(224,230,235,0));
+        }
+        .group-divider .group-label {
+            display: inline-block;
+            border: 1px solid var(--accent-muted);
+            border-radius: 999px;
+            padding: 2px 10px;
+            font-size: var(--group-label-font-size);
+            font-weight: 700;
+            color: var(--muted-text);
+            background: var(--bg);
+            box-shadow: var(--inset-shadow);
+        }
+        .chat-turn .turn-content { position: relative; }
+        .chat-turn .turn-content::before {
+            content: none;
+        }
+        .chat-turn .agent-badge {
+            position: relative;
+        }
+        /* removed absolute-positioned emoji to prevent overlap */
+    </style>
+    """
+
+    # HTML structure
+    html_parts = [
+        "<!DOCTYPE html>",
+        "<html>",
+        "<head>",
+        "<meta charset='UTF-8'>",
+        "<title>Chat Turns</title>",
+        css,
+        "<script>\n"
+        "document.addEventListener('DOMContentLoaded', function() {\n"
+        "  const flow = document.querySelector('.messages-flow');\n"
+        "  // State for range filtering and strong hide\n"
+        "  let currentRangeStart = null;\n"
+        "  let currentRangeEnd = null;\n"
+        "  let strongHideOn = false;\n"
+        "  // Toggle collapse per message\n"
+        "  document.body.addEventListener('click', function(e){\n"
+        "    if (e.target.closest('.ts-badge')) { return; }\n"
+        "    if (e.target.closest('.agent-badge')) { return; }\n"
+        "    const seg = e.target.closest('.segment');\n"
+        "    if (seg) { e.stopPropagation(); seg.classList.toggle('collapsed'); }\n"
+        "  });\n"
+        "  // Grouping logic\n"
+        "  function applyRangeFilter() {\n"
+        "    const turns = Array.from(flow.querySelectorAll('.chat-turn'));\n"
+        "    for (const el of turns) {\n"
+        "      const t = parseInt(el.getAttribute('data-time-step') || '0', 10);\n"
+        "      const afterStart = (currentRangeStart === null) || (t >= currentRangeStart);\n"
+        "      const beforeEnd = (currentRangeEnd === null) || (t <= currentRangeEnd);\n"
+        "      el.style.display = (afterStart && beforeEnd) ? '' : 'none';\n"
+        "    }\n"
+        "    // Hide group headers that have no visible turns in their section\n"
+        "    const dividers = Array.from(flow.querySelectorAll('.group-divider'));\n"
+        "    for (const d of dividers) {\n"
+        "      let anyVisible = false;\n"
+        "      let el = d.nextElementSibling;\n"
+        "      while (el && !el.classList.contains('group-divider')) {\n"
+        "        if (el.classList.contains('chat-turn')) {\n"
+        "          const disp = getComputedStyle(el).display;\n"
+        "          if (disp !== 'none') { anyVisible = true; break; }\n"
+        "        }\n"
+        "        el = el.nextElementSibling;\n"
+        "      }\n"
+        "      d.style.display = anyVisible ? '' : 'none';\n"
+        "    }\n"
+        "  }\n"
+        "  function applyGrouping(n) {\n"
+        "    // Remove existing group dividers\n"
+        "    Array.from(flow.querySelectorAll('.group-divider')).forEach(el => el.remove());\n"
+        "    if (!n || n <= 0) { return; }\n"
+        "    const turns = Array.from(flow.querySelectorAll('.chat-turn'));\n"
+        "    if (turns.length === 0) return;\n"
+        "    // Re-append in order with dividers\n"
+        "    const items = Array.from(flow.children).filter(el => !el.classList.contains('group-divider'));\n"
+        "    const frag = document.createDocumentFragment();\n"
+        "    let lastGroup = -1;\n"
+        "    for (const el of items) {\n"
+        "      if (!el.classList.contains('chat-turn')) { frag.appendChild(el); continue; }\n"
+        "      const t = parseInt(el.getAttribute('data-time-step') || '0', 10);\n"
+        "      const g = Math.floor(t / n);\n"
+        "      if (g !== lastGroup) {\n"
+        "        const div = document.createElement('div');\n"
+        "        div.className = 'group-divider';\n"
+        "        const label = document.createElement('span');\n"
+        "        label.className = 'group-label';\n"
+        "        const start = g * n;\n"
+        "        const end = start + n - 1;\n"
+        "        const roundIndex = g + 1;\n"
+        "        label.textContent = `Round ${roundIndex}`;\n"
+        "        div.appendChild(label);\n"
+        "        frag.appendChild(div);\n"
+        "        lastGroup = g;\n"
+        "      }\n"
+        "      frag.appendChild(el);\n"
+        "    }\n"
+        "    flow.innerHTML = '';\n"
+        "    flow.appendChild(frag);\n"
+        "    // Hide timestep badges when grouping is 1\n"
+        "    flow.classList.toggle('hide-ts-badges', n === 1);\n"
+        "    // Keep strong hide state\n"
+        "    flow.classList.toggle('strong-hide', strongHideOn);\n"
+        "    // Re-apply range filter after regrouping\n"
+        "    applyRangeFilter();\n"
+        "  }\n"
+        "  const input = document.getElementById('group-size');\n"
+        "  const btn = document.getElementById('apply-grouping');\n"
+        "  if (btn && input) {\n"
+        "    btn.addEventListener('click', () => { const n = parseInt(input.value || '0', 10); applyGrouping(n); });\n"
+        "    input.addEventListener('keydown', (e) => { if (e.key === 'Enter') { const n = parseInt(input.value || '0', 10); applyGrouping(n); } });\n"
+        "  }\n"
+        "  // Default grouping to 1 timestep on load\n"
+        "  if (input) { input.value = '1'; applyGrouping(1); }\n"
+        "  // Range filter controls\n"
+        "  const rangeStart = document.getElementById('range-start');\n"
+        "  const rangeEnd = document.getElementById('range-end');\n"
+        "  const rangeBtn = document.getElementById('apply-range');\n"
+        "  if (rangeBtn && rangeStart && rangeEnd) {\n"
+        "    const applyRange = () => {\n"
+        "      const sv = parseInt(rangeStart.value || '', 10);\n"
+        "      const ev = parseInt(rangeEnd.value || '', 10);\n"
+        "      currentRangeStart = Number.isFinite(sv) ? sv : null;\n"
+        "      currentRangeEnd = Number.isFinite(ev) ? ev : null;\n"
+        "      applyRangeFilter();\n"
+        "    };\n"
+        "    rangeBtn.addEventListener('click', applyRange);\n"
+        "    rangeStart.addEventListener('keydown', (e) => { if (e.key === 'Enter') applyRange(); });\n"
+        "    rangeEnd.addEventListener('keydown', (e) => { if (e.key === 'Enter') applyRange(); });\n"
+        "  }\n"
+        "  // Strong hide toggle (on by default)\n"
+        "  const strongHideBtn = document.getElementById('toggle-strong-hide');\n"
+        "  const strongHideStateEl = document.getElementById('strong-hide-state');\n"
+        "  if (strongHideBtn) {\n"
+        "    const setLabel = () => { if (strongHideStateEl) { strongHideStateEl.textContent = strongHideOn ? 'On' : 'Off'; } };\n"
+        "    strongHideBtn.addEventListener('click', () => { strongHideOn = !strongHideOn; flow.classList.toggle('strong-hide', strongHideOn); setLabel(); });\n"
+        "    flow.classList.toggle('strong-hide', strongHideOn);\n"
+        "    setLabel();\n"
+        "  }\n"
+        "});\n"
+        "</script>",
+        "</head>",
+        "<body>",
+        '<div class="toolbar-wrap">',
+        '<div class="toolbar-hotzone"></div>',
+        '<div class="toolbar">',
+        '<label for="group-size">Group every</label>',
+        '<input id="group-size" type="number" min="0" step="1" value="1" />',
+        "<span>timesteps</span>",
+        '<button id="apply-grouping">Apply</button>',
+        '<span style="margin-left:8px"></span>',
+        '<label for="range-start"><span class="emoji-bw">🔎</span> Range</label>',
+        '<input id="range-start" type="number" step="1" />',
+        "<span>to</span>",
+        '<input id="range-end" type="number" step="1" />',
+        '<button id="apply-range"><span class="emoji-bw">▶︎</span> Apply</button>',
+        '<button id="toggle-strong-hide"><span class="emoji-bw">🗜️</span> Strong Hide: <span id="strong-hide-state">On</span></button>',
+        "</div>",
+        "</div>",
+        '<div class="messages-flow">',
+    ]
+
+    last_time_step = None
+    for original_index, turn in indexed_turns:
+        # Build classes
+        agent_class = f"agent-{re.sub('[^a-z0-9_-]', '-', turn.agent_id.lower())}"
+        role_class = f"role-{turn.role}"
+        # Segments default collapsed for user role
+        segment_collapsed_class = " collapsed" if turn.role == "user" else ""
+
+        # Badge content
+        if turn.role == "assistant":
+            name = html.escape(turn.agent_id)
+            emoji = '<span class="emoji-bw">🤖</span>'
+            raw_val = turn.reward
+            if isinstance(raw_val, (int, float)):
+                reward_val = f"{raw_val:.4f}".rstrip("0").rstrip(".")
+                if len(reward_val) > 8:
+                    reward_val = reward_val[:8] + "…"
+            else:
+                reward_val = str(raw_val)
+            # Format: "🤖 Alice 💬 • Reward: 5.5556 • "
+            badge_inner = (
+                f'{emoji} <span class="agent-name">{name}</span>'
+                f' <span class="sep"> • </span><span class="reward">{reward_val} r</span>'
+                f' <span class="sep"> • </span>'
+            )
+        else:
+            # For user messages, show "User of {Agent ID}" in the badge
+            name = "User of " + html.escape(turn.agent_id)
+            emoji = '<span class="emoji-bw">⚙️</span>'
+            # Format (no reward): "⚙️ User of Alice • "
+            badge_inner = f'{emoji} <span class="agent-name">{name}</span><span class="sep"> • </span>'
+
+        badge = f'<span class="agent-badge">{badge_inner}</span>'
+
+        # Inline timestep distinction badge at step boundaries (render before first message)
+        ts_badge_html = ""
+        if last_time_step is None or turn.time_step != last_time_step:
+            ts_badge_html = f'<span class="ts-badge">⏱ {turn.time_step}</span>'
+            last_time_step = turn.time_step
+
+        escaped_content = html.escape(turn.content)
+        collapsed_text = re.sub(r"\s+", " ", escaped_content).strip()
+        # Optional reasoning
+        reasoning_val = getattr(turn, "reasoning", None)
+        reasoning_html = ""
+        if reasoning_val:
+            escaped_reasoning = html.escape(reasoning_val)
+            reasoning_text = re.sub(r"\s+", " ", escaped_reasoning).strip()
+            reasoning_html = (
+                f'<span class="segment reasoning-box collapsed{segment_collapsed_class}">'
+                f'<span class="emoji-bw">💭</span> '
+                f'<span class="seg-text"><i>{reasoning_text}  </i></span>'
+                f"</span>"
+            )
+
+        html_parts.append(
+            f'<div class="chat-turn {agent_class} {role_class}" data-time-step="{turn.time_step}">'
+            f'<div class="turn-content {agent_class} {role_class}">{ts_badge_html}{badge}'
+            f"{reasoning_html}"
+            f'<span class="segment message-box{segment_collapsed_class}"><span class="emoji-bw">💬</span> <span class="seg-text">{collapsed_text}</span></span>'
+            f"</div>"
+            f"</div>"
+        )
+
+    html_parts.extend(["</div>", "</body>", "</html>"])
+
+    return "\n".join(html_parts)
+
+
+def export_html_from_rollout_tree(path: Path, outdir: Path, main_only: bool = False):
+    """Process a rollout tree file and generate HTML files for each path.
+    Creates separate HTML files for the main path and each branch path.
+    The main path is saved in the root output directory, while branch paths
+    are saved in a 'branches' subdirectory.
+
+    Args:
+        path: Path to the rollout tree JSON file
+        outdir: Output directory for HTML files
+        main_only: If True, only export the main trajectory (default: False)
+    """
+    root = load_rollout_tree(path)
+    mgid = root.id
+
+    main_path, branch_paths = get_rollout_tree_paths(root)
+
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    # Create branches subdirectory if we have branch paths
+    if not main_only and branch_paths:
+        branches_dir = outdir / f"mgid:{mgid}_branches_html_renders"
+        branches_dir.mkdir(parents=True, exist_ok=True)
+
+    # Generate HTML for the main path
+    chat_turns = gather_all_chat_turns_for_path(main_path)
+    html_content = html_from_chat_turns(chat_turns)
+    output_file = outdir / f"mgid:{mgid}_main_html_render.render.html"
+    with open(output_file, "w", encoding="utf-8") as f:
+        f.write(html_content)
+
+    # Generate HTML for each branch path
+    for path_obj in branch_paths:
+        chat_turns = gather_all_chat_turns_for_path(path_obj)
+
+        html_content = html_from_chat_turns(chat_turns)
+
+        path_id: str = path_obj.id
+        output_filename = f"{path_id}_html_render.render.html"
+
+        output_file = branches_dir / output_filename
+
+        with open(output_file, "w", encoding="utf-8") as f:
+            f.write(html_content)

src_code_for_reproducibility/markov_games/mg_utils.py

@@ -0,0 +1,77 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+import copy
+import asyncio
+
+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.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_simulation import TrustAndSplitSimulation
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import RolloutTreeRootNode, StepLog, RolloutTreeBranchNode
+from mllm.markov_games.rollout_tree import AgentActLog
+from mllm.markov_games.simulation import SimulationStepLog
+from mllm.markov_games.rollout_tree import RolloutTreeNode
+
+AgentId = str
+
+
+
+
+@dataclass
+class AgentConfig:
+    agent_id: int
+    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
+    """
+    simulation = eval(config.simulation_class_name)(
+        seed=config.seed,
+        **config.simulation_init_args,
+    )
+    agents = {}
+    for agent_config in config.agent_configs:
+        agent_id = agent_config.agent_id
+        agent_class = eval(agent_config.agent_class_name)
+        agent = agent_class(
+            seed=config.seed,
+            agent_id=agent_id,
+            policy=policies[agent_config.policy_id],
+            **agent_config.init_kwargs,
+        )
+        agents[agent_id] = agent
+    markov_game = MarkovGame(
+        id=config.id,
+        crn_id=config.seed,
+        simulation=simulation,
+        agents=agents,
+    )
+    return markov_game
+
+
+
+

src_code_for_reproducibility/markov_games/simulation.py

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

src_code_for_reproducibility/markov_games/statistics_runner.py

@@ -0,0 +1,405 @@
+from __future__ import annotations
+
+import gc
+import json
+import pickle
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional
+
+from basic_render import find_iteration_folders
+
+from mllm.markov_games.rollout_tree import (
+    RolloutTreeBranchNode,
+    RolloutTreeNode,
+    RolloutTreeRootNode,
+    SimulationStepLog,
+)
+
+
+def _iterate_main_nodes(root: RolloutTreeRootNode) -> Iterator[RolloutTreeNode]:
+    """
+    Iterate the main path nodes without materializing full path lists.
+    """
+    current = root.child
+    while current is not None:
+        if isinstance(current, RolloutTreeNode):
+            yield current
+            current = current.child
+        elif isinstance(current, RolloutTreeBranchNode):
+            # Follow only the main child on the main trajectory
+            current = current.main_child
+        else:
+            break
+
+
+def iterate_main_simulation_logs(
+    root: RolloutTreeRootNode,
+) -> Iterator[SimulationStepLog]:
+    for node in _iterate_main_nodes(root):
+        yield node.step_log.simulation_step_log
+
+
+def stream_rollout_files(iteration_folder: Path) -> Iterator[Path]:
+    for p in iteration_folder.rglob("*.rt.pkl"):
+        if p.is_file():
+            yield p
+
+
+def load_root(path: Path) -> RolloutTreeRootNode:
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+    return RolloutTreeRootNode.model_validate(data)
+
+
+@dataclass
+class StatRecord:
+    mgid: int
+    crn_id: Optional[int]
+    iteration: str
+    values: Dict[str, Any]
+
+
+class StatComputer:
+    """
+    Stateful stat computer that consumes SimulationStepLog instances
+    and produces final aggregated values for one rollout (mgid).
+    """
+
+    def update(self, sl: SimulationStepLog) -> None:  # pragma: no cover - interface
+        raise NotImplementedError
+
+    def finalize(self) -> Dict[str, Any]:  # pragma: no cover - interface
+        raise NotImplementedError
+
+
+def run_stats(
+    data_root: Path,
+    game_name: str,
+    make_computers: Callable[[], List[StatComputer]],
+    output_filename: Optional[str] = None,
+    output_format: str = "json",  # "json" (dict of lists) or "jsonl"
+) -> Path:
+    """
+    Compute stats across all iteration_* folders under data_root.
+    Writes JSONL to data_root/statistics/<output_filename or f"{game_name}.stats.jsonl">.
+    """
+    data_root = Path(data_root)
+    outdir = data_root / "statistics"
+    outdir.mkdir(parents=True, exist_ok=True)
+    # Choose extension by format
+    default_name = (
+        f"{game_name}.stats.json"
+        if output_format == "json"
+        else f"{game_name}.stats.jsonl"
+    )
+    outfile = outdir / (
+        output_filename if output_filename is not None else default_name
+    )
+
+    # Rewrite file each run to keep it clean and small
+    if outfile.exists():
+        outfile.unlink()
+
+    iteration_folders = find_iteration_folders(str(data_root))
+
+    # If writing JSONL, stream directly; otherwise accumulate minimal records
+    if output_format == "jsonl":
+        with open(outfile, "w", encoding="utf-8") as w:
+            for iteration_folder in iteration_folders:
+                iteration_name = Path(iteration_folder).name
+                for pkl_path in stream_rollout_files(Path(iteration_folder)):
+                    root = load_root(pkl_path)
+
+                    computers = make_computers()
+                    for sl in iterate_main_simulation_logs(root):
+                        for comp in computers:
+                            try:
+                                comp.update(sl)
+                            except Exception:
+                                continue
+
+                    values: Dict[str, Any] = {}
+                    for comp in computers:
+                        try:
+                            values.update(comp.finalize())
+                        except Exception:
+                            continue
+
+                    rec = {
+                        "mgid": getattr(root, "id", None),
+                        "crn_id": getattr(root, "crn_id", None),
+                        "iteration": iteration_name,
+                        "stats": values,
+                    }
+                    w.write(json.dumps(rec, ensure_ascii=False) + "\n")
+
+                    del root
+                    del computers
+                    gc.collect()
+    else:
+        # Aggregate to dict-of-lists for easier plotting
+        records: List[Dict[str, Any]] = []
+        # Process in deterministic order
+        for iteration_folder in iteration_folders:
+            iteration_name = Path(iteration_folder).name
+            for pkl_path in stream_rollout_files(Path(iteration_folder)):
+                root = load_root(pkl_path)
+
+                computers = make_computers()
+                for sl in iterate_main_simulation_logs(root):
+                    for comp in computers:
+                        try:
+                            comp.update(sl)
+                        except Exception:
+                            continue
+
+                values: Dict[str, Any] = {}
+                for comp in computers:
+                    try:
+                        values.update(comp.finalize())
+                    except Exception:
+                        continue
+
+                records.append(
+                    {
+                        "mgid": getattr(root, "id", None),
+                        "crn_id": getattr(root, "crn_id", None),
+                        "iteration": iteration_name,
+                        "stats": values,
+                    }
+                )
+
+                del root
+                del computers
+                gc.collect()
+
+        # Build dict-of-lists with nested stats preserved
+        # Collect all stat keys and nested agent keys where needed
+        mgids: List[Any] = []
+        crn_ids: List[Any] = []
+        iterations_out: List[str] = []
+        # stats_out is a nested structure mirroring keys but with lists
+        stats_out: Dict[str, Any] = {}
+
+        # First pass to collect union of keys
+        stat_keys: set[str] = set()
+        nested_agent_keys: Dict[str, set[str]] = {}
+        for r in records:
+            stats = r.get("stats", {}) or {}
+            for k, v in stats.items():
+                stat_keys.add(k)
+                if isinstance(v, dict):
+                    nested = nested_agent_keys.setdefault(k, set())
+                    for ak in v.keys():
+                        nested.add(str(ak))
+
+        # Initialize structure
+        for k in stat_keys:
+            if k in nested_agent_keys:
+                stats_out[k] = {ak: [] for ak in sorted(nested_agent_keys[k])}
+            else:
+                stats_out[k] = []
+
+        # Fill lists
+        for r in records:
+            mgids.append(r.get("mgid"))
+            crn_ids.append(r.get("crn_id"))
+            iterations_out.append(r.get("iteration"))
+            stats = r.get("stats", {}) or {}
+            for k in stat_keys:
+                val = stats.get(k)
+                if isinstance(stats_out[k], dict):
+                    # per-agent dict
+                    agent_dict = val if isinstance(val, dict) else {}
+                    for ak in stats_out[k].keys():
+                        stats_out[k][ak].append(agent_dict.get(ak))
+                else:
+                    stats_out[k].append(val)
+
+        with open(outfile, "w", encoding="utf-8") as w:
+            json.dump(
+                {
+                    "mgid": mgids,
+                    "crn_id": crn_ids,
+                    "iteration": iterations_out,
+                    "stats": stats_out,
+                },
+                w,
+                ensure_ascii=False,
+            )
+
+    return outfile
+
+
+def run_stats_functional(
+    data_root: Path,
+    game_name: str,
+    metrics: Dict[str, Callable[[SimulationStepLog], Optional[Dict[str, float]]]],
+    output_filename: Optional[str] = None,
+    output_format: str = "json",
+) -> Path:
+    """
+    Functional variant where metrics is a dict of name -> f(SimulationStepLog) -> {agent_id: value}.
+    Aggregates per rollout by averaging over steps where a metric produced a value.
+    Writes a single consolidated file in data_root/statistics/.
+    """
+    data_root = Path(data_root)
+    outdir = data_root / "statistics"
+    outdir.mkdir(parents=True, exist_ok=True)
+    default_name = (
+        f"{game_name}.stats.json"
+        if output_format == "json"
+        else f"{game_name}.stats.jsonl"
+    )
+    outfile = outdir / (
+        output_filename if output_filename is not None else default_name
+    )
+
+    if outfile.exists():
+        outfile.unlink()
+
+    iteration_folders = find_iteration_folders(str(data_root))
+
+    def finalize_rollout(
+        agg: Dict[str, Dict[str, List[float]]]
+    ) -> Dict[str, Dict[str, float]]:
+        # avg per metric per agent
+        result: Dict[str, Dict[str, float]] = {}
+        for mname, agent_values in agg.items():
+            result[mname] = {}
+            for aid, vals in agent_values.items():
+                if not vals:
+                    result[mname][aid] = None  # keep alignment; could be None
+                else:
+                    result[mname][aid] = sum(vals) / len(vals)
+        return result
+
+    if output_format == "jsonl":
+        with open(outfile, "w", encoding="utf-8") as w:
+            for iteration_folder in iteration_folders:
+                iteration_name = Path(iteration_folder).name
+                for pkl_path in stream_rollout_files(Path(iteration_folder)):
+                    root = load_root(pkl_path)
+
+                    # aggregator structure: metric -> agent_id -> list of values
+                    agg: Dict[str, Dict[str, List[float]]] = {
+                        m: {} for m in metrics.keys()
+                    }
+
+                    for sl in iterate_main_simulation_logs(root):
+                        for mname, fn in metrics.items():
+                            try:
+                                vals = fn(sl)
+                            except Exception:
+                                vals = None
+                            if not vals:
+                                continue
+                            for aid, v in vals.items():
+                                if v is None:
+                                    continue
+                                lst = agg[mname].setdefault(str(aid), [])
+                                try:
+                                    lst.append(float(v))
+                                except Exception:
+                                    continue
+
+                    values = finalize_rollout(agg)
+                    rec = {
+                        "mgid": getattr(root, "id", None),
+                        "crn_id": getattr(root, "crn_id", None),
+                        "iteration": iteration_name,
+                        "stats": values,
+                    }
+                    w.write(json.dumps(rec, ensure_ascii=False) + "\n")
+
+                    del root
+                    gc.collect()
+    else:
+        records: List[Dict[str, Any]] = []
+        for iteration_folder in iteration_folders:
+            iteration_name = Path(iteration_folder).name
+            for pkl_path in stream_rollout_files(Path(iteration_folder)):
+                root = load_root(pkl_path)
+
+                agg: Dict[str, Dict[str, List[float]]] = {m: {} for m in metrics.keys()}
+                for sl in iterate_main_simulation_logs(root):
+                    for mname, fn in metrics.items():
+                        try:
+                            vals = fn(sl)
+                        except Exception:
+                            vals = None
+                        if not vals:
+                            continue
+                        for aid, v in vals.items():
+                            if v is None:
+                                continue
+                            lst = agg[mname].setdefault(str(aid), [])
+                            try:
+                                lst.append(float(v))
+                            except Exception:
+                                continue
+
+                values = finalize_rollout(agg)
+                records.append(
+                    {
+                        "mgid": getattr(root, "id", None),
+                        "crn_id": getattr(root, "crn_id", None),
+                        "iteration": iteration_name,
+                        "stats": values,
+                    }
+                )
+
+                del root
+                gc.collect()
+
+        # Build dict-of-lists output
+        mgids: List[Any] = []
+        crn_ids: List[Any] = []
+        iterations_out: List[str] = []
+        stats_out: Dict[str, Any] = {}
+
+        stat_keys: set[str] = set()
+        nested_agent_keys: Dict[str, set[str]] = {}
+        for r in records:
+            stats = r.get("stats", {}) or {}
+            for k, v in stats.items():
+                stat_keys.add(k)
+                if isinstance(v, dict):
+                    nested = nested_agent_keys.setdefault(k, set())
+                    for ak in v.keys():
+                        nested.add(str(ak))
+
+        for k in stat_keys:
+            if k in nested_agent_keys:
+                stats_out[k] = {ak: [] for ak in sorted(nested_agent_keys[k])}
+            else:
+                stats_out[k] = []
+
+        for r in records:
+            mgids.append(r.get("mgid"))
+            crn_ids.append(r.get("crn_id"))
+            iterations_out.append(r.get("iteration"))
+            stats = r.get("stats", {}) or {}
+            for k in stat_keys:
+                val = stats.get(k)
+                if isinstance(stats_out[k], dict):
+                    agent_dict = val if isinstance(val, dict) else {}
+                    for ak in stats_out[k].keys():
+                        stats_out[k][ak].append(agent_dict.get(ak))
+                else:
+                    stats_out[k].append(val)
+
+        with open(outfile, "w", encoding="utf-8") as w:
+            json.dump(
+                {
+                    "mgid": mgids,
+                    "crn_id": crn_ids,
+                    "iteration": iterations_out,
+                    "stats": stats_out,
+                },
+                w,
+                ensure_ascii=False,
+            )
+
+    return outfile

src_code_for_reproducibility/models/adapter_training_wrapper.py

@@ -0,0 +1,89 @@
+
+import torch.nn as nn
+import logging
+from typing import Union
+from peft import (
+    LoraConfig,
+    get_peft_model,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AdapterWrapper(nn.Module):
+    """
+    A thin façade that
+      • keeps a reference to a *shared* PEFT-wrapped model,
+      • ensures `set_adapter(adapter)` is called on every forward,
+      • exposes only the parameters that should be trained for that adapter
+        (plus whatever extra modules you name).
+    """
+    def __init__(
+        self,
+        shared_llm: nn.Module,
+        adapter_id: str,
+        lora_config: dict,
+        path: Union[str, None] = None,
+        ):
+        super().__init__()
+        self.shared_llm = shared_llm
+        self.adapter_id = adapter_id
+        lora_config = LoraConfig(**lora_config)
+        # this modifies the shared llm in place, adding a lora adapter inside
+        self.shared_llm = get_peft_model(
+            model=shared_llm,
+            peft_config=lora_config,
+            adapter_name=adapter_id,
+        )
+        self.shared_llm.train()
+        # Load external adapter weights if provided
+        loaded_from: str | None = None
+        if path:
+            try:
+                # Supports both local filesystem paths and HF Hub repo IDs
+                self.shared_llm.load_adapter(
+                    is_trainable=True,
+                    model_id=path,
+                    adapter_name=adapter_id,
+                )
+                loaded_from = path
+            except Exception as exc:  # noqa: BLE001 - want to log any load failure context
+                logger.warning(
+                    f"Adapter '{adapter_id}': failed to load from '{path}': {exc}"
+                )
+
+        if loaded_from:
+            logger.info(
+                f"Adapter '{adapter_id}': loaded initial weights from '{loaded_from}'."
+            )
+        else:
+            logger.info(
+                f"Adapter '{adapter_id}': initialized with fresh weights (no initial weights found)."
+            )
+
+    def parameters(self, recurse: bool = True):
+        """
+        "recurse" is just for pytorch compatibility
+        """
+        self.shared_llm.set_adapter(self.adapter_id)
+        params = [p for p in self.shared_llm.parameters() if p.requires_grad]
+
+        return params
+
+    def forward(self, *args, **kwargs):
+        self.shared_llm.set_adapter(self.adapter_id)
+        return self.shared_llm(*args, **kwargs)
+
+    def save_pretrained(self, save_path):
+        self.shared_llm.save_pretrained(save_path)
+
+    def gradient_checkpointing_enable(self, *args, **kwargs):
+        self.shared_llm.gradient_checkpointing_enable(*args, **kwargs)
+
+    @property
+    def dtype(self):
+        return self.shared_llm.dtype
+
+    @property
+    def device(self):
+        return self.shared_llm.device

src_code_for_reproducibility/models/inference_backend.py

@@ -0,0 +1,35 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+from dataclasses import dataclass
+
+@dataclass
+class PolicyOutput:
+    content: str
+    reasoning_content: str | None = None
+
+class LLMInferenceBackend(ABC):
+    @abstractmethod
+    def __init__(self, **kwargs):
+        ...
+
+    @abstractmethod
+    def prepare_adapter(
+        self, adapter_id: str, weights_got_updated: bool = False
+    ) -> None:
+        """Ensure adapter is ready/loaded for next generation call."""
+
+    @abstractmethod
+    async def generate(self, prompt: list[dict], regex: Optional[str] = None) -> PolicyOutput:
+        ...
+
+    @abstractmethod
+    def toggle_training_mode(self) -> None:
+        ...
+
+    @abstractmethod
+    def toggle_eval_mode(self) -> None:
+        ...
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        ...

src_code_for_reproducibility/models/inference_backend_dummy.py

@@ -0,0 +1,53 @@
+import asyncio
+import re
+from typing import Optional
+
+import rstr
+from transformers import AutoTokenizer
+
+from mllm.models.inference_backend import LLMInferenceBackend, PolicyOutput
+from mllm.utils.short_id_gen import generate_short_id
+
+
+class DummyInferenceBackend(LLMInferenceBackend):
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        pass
+
+    def prepare_adapter(
+        self, adapter_id: Optional[str], weights_got_updated: bool
+    ) -> None:
+        pass
+
+    async def toggle_training_mode(self) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def toggle_eval_mode(self) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    async def generate(
+        self, prompt_text: str, regex: Optional[str] = None
+    ) -> PolicyOutput:
+        content = "I am a dummy backend without a regex."
+        reasoning_content = None
+
+        if regex:
+            raw_text = rstr.xeger(regex)
+            content = raw_text
+            # Strict split: require \n<think>...</think>\n\n before final content
+            m = re.match(
+                r"^\n<think>\n([\s\S]*?)</think>\n\n(.*)$", raw_text, flags=re.DOTALL
+            )
+            if m:
+                reasoning_content = m.group(1)
+                content = m.group(2)
+
+        return PolicyOutput(content=content, reasoning_content=reasoning_content)

src_code_for_reproducibility/models/inference_backend_sglang.py

@@ -0,0 +1,86 @@
+# new_backend_sglang_offline.py
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Optional
+
+import sglang as sgl
+
+from mllm.models.inference_backend import LLMInferenceBackend
+
+
+class SGLangOfflineBackend(LLMInferenceBackend):
+    def __init__(
+        self,
+        model_name: str,
+        tokenizer,  # unused but kept for parity
+        adapter_paths: dict[str, str],
+        device: str = "cuda",
+        max_model_len: Optional[int] = None,
+        enable_lora: bool = True,
+        lora_target_modules: Optional[list[str] | str] = None,
+        max_loras_per_batch: int = 8,
+        engine_kwargs: dict[str, Any] = None,
+    ):
+        self.model_name = model_name
+        self.adapter_paths = adapter_paths
+        self.current_adapter: Optional[str] = None
+        engine_kwargs = dict(engine_kwargs or {})
+        # Map server-style LoRA flags to offline engine ctor
+        if enable_lora and adapter_paths:
+            engine_kwargs.setdefault("enable_lora", True)
+            # The offline Engine mirrors server args; pass a mapping name->path
+            engine_kwargs.setdefault("lora_paths", adapter_paths)
+            if lora_target_modules is not None:
+                engine_kwargs.setdefault("lora_target_modules", lora_target_modules)
+            engine_kwargs.setdefault("max_loras_per_batch", max_loras_per_batch)
+
+        if max_model_len is not None:
+            engine_kwargs.setdefault("context_length", max_model_len)
+
+        # Launch in-process engine (no HTTP server)
+        self.llm = sgl.Engine(model_path=model_name, **engine_kwargs)  # async-ready
+        # SGLang supports: generate(), async_generate(), and async streaming helpers. :contentReference[oaicite:2]{index=2}
+
+    def is_ready(self) -> bool:
+        return True
+
+    def toggle_training_mode(self) -> None:
+        # No explicit KV release API offline; typically you pause usage here.
+        pass
+
+    def toggle_eval_mode(self) -> None:
+        pass
+
+    def shutdown(self) -> None:
+        # Engine cleans up on GC; explicit close not required.
+        pass
+
+    def prepare_adapter(self, adapter_id: Optional[str]) -> None:
+        # With offline Engine, when LoRA is enabled at init,
+        # you select adapter per request via the input batch mapping.
+        self.current_adapter = adapter_id
+
+    async def generate(
+        self, prompt_text: str, sampling_params: dict, adapter_id: Optional[str]
+    ) -> str:
+        # Non-streaming async (batch of 1). For batched prompts, pass a list.
+        params = {
+            "temperature": sampling_params.get("temperature", 1.0),
+            "top_p": sampling_params.get("top_p", 1.0),
+            "max_new_tokens": sampling_params.get("max_new_tokens", 128),
+        }
+        if (tk := sampling_params.get("top_k", -1)) and tk > 0:
+            params["top_k"] = tk
+        if (mn := sampling_params.get("min_new_tokens")) is not None:
+            params["min_new_tokens"] = mn
+        if (fp := sampling_params.get("frequency_penalty")) is not None:
+            params["frequency_penalty"] = fp
+
+        # If using multi-LoRA, SGLang lets you provide adapter names aligned to each input.
+        prompts = [prompt_text]
+        adapters = [adapter_id] if adapter_id else None  # or omit for base
+        outs = await self.llm.async_generate(
+            prompts, params, adapters
+        )  # :contentReference[oaicite:3]{index=3}
+        return outs[0]["text"]

src_code_for_reproducibility/models/inference_backend_sglang_local_server.py

@@ -0,0 +1,127 @@
+import os
+
+import httpx
+import requests
+from sglang.utils import launch_server_cmd, wait_for_server
+
+from mllm.models.inference_backend import LLMInferenceBackend
+
+
+class HttpSGLangBackend(LLMInferenceBackend):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.port = None
+        self.proc = None
+        self.urls = {}
+        # track sglang adapter ids separately from your logical ids
+        self.sglang_names = {aid: aid for aid in self.adapter_paths.keys()}
+        self.needs_loading = {aid: True for aid in self.adapter_paths.keys()}
+
+        # defaults you already used:
+        self.mem_fraction = kwargs.get("mem_fraction_static", 0.6)
+        self.dtype = kwargs.get("dtype", "bfloat16")
+        self.extra_cli = kwargs.get("extra_cli", "")
+        self.disable_radix_cache = kwargs.get("disable_radix_cache", True)
+
+    def launch(self) -> None:
+        # find local hf cache path for server
+        from transformers.utils import cached_file
+
+        local_llm_path = os.path.split(cached_file(self.model_name, "config.json"))[0]
+
+        lora_str = ""
+        if self.adapter_paths:
+            lora_str = "--lora-paths " + " ".join(
+                f"{aid}={path}" for aid, path in self.adapter_paths.items()
+            )
+
+        cmd = f"""
+        python3 -m sglang.launch_server --model-path {local_llm_path} \
+        --host 0.0.0.0 {lora_str} \
+        {'--disable-radix-cache' if self.disable_radix_cache else ''} \
+        --mem-fraction-static {self.mem_fraction} --dtype {self.dtype} {self.extra_cli}
+        """
+        self.proc, self.port = launch_server_cmd(cmd)
+        wait_for_server(f"http://localhost:{self.port}")
+        base = f"http://localhost:{self.port}"
+        self.urls = dict(
+            generate=f"{base}/generate",
+            release=f"{base}/release_memory_occupation",
+            resume=f"{base}/resume_memory_occupation",
+            load_lora=f"{base}/load_lora_adapter",
+            unload_lora=f"{base}/unload_lora_adapter",
+        )
+
+    def is_ready(self) -> bool:
+        try:
+            requests.get(self.urls["generate"], timeout=2)
+            return True
+        except Exception:
+            return False
+
+    def prepare_adapter(self, adapter_id: str) -> None:
+        if adapter_id is None:
+            return
+        if self.needs_loading.get(adapter_id, False):
+            # unload old name if present
+            try:
+                requests.post(
+                    self.urls["unload_lora"],
+                    json={"lora_name": self.sglang_names[adapter_id]},
+                    timeout=10,
+                )
+            except Exception:
+                pass
+            new_name = self._short_id()
+            self.sglang_names[adapter_id] = new_name
+            requests.post(
+                self.urls["load_lora"],
+                json={
+                    "lora_name": new_name,
+                    "lora_path": self.adapter_paths[adapter_id],
+                },
+            ).raise_for_status()
+            self.needs_loading[adapter_id] = False
+
+    async def generate(
+        self, prompt_text: str, sampling_params: dict, adapter_id: str | None
+    ) -> str:
+        lora_name = self.sglang_names.get(adapter_id) if adapter_id else None
+        payload = {
+            "text": [prompt_text],
+            "sampling_params": sampling_params,
+        }
+        if lora_name:
+            payload["lora_path"] = [lora_name]
+
+        timeout = httpx.Timeout(3600.0, connect=3600.0)
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            resp = await client.post(self.urls["generate"], json=payload)
+            resp.raise_for_status()
+            return resp.json()[0]["text"]
+
+    def toggle_training_mode(self) -> None:
+        # free KV space while training adapters
+        requests.post(
+            self.urls["release"], json={"tags": ["kv_cache"]}
+        ).raise_for_status()
+
+    def toggle_eval_mode(self) -> None:
+        # re-allocate KV space
+        try:
+            requests.post(
+                self.urls["resume"], json={"tags": ["kv_cache"]}
+            ).raise_for_status()
+        except Exception:
+            pass
+
+    def shutdown(self) -> None:
+        from sglang.utils import terminate_process
+
+        if self.proc:
+            terminate_process(self.proc)
+
+    def _short_id(self) -> str:
+        import uuid
+
+        return str(uuid.uuid4().int)[:8]

src_code_for_reproducibility/models/inference_backend_vllm.py

@@ -0,0 +1,96 @@
+import asyncio
+import re
+from typing import Optional
+
+from transformers import AutoTokenizer
+from vllm import AsyncEngineArgs, AsyncLLMEngine, SamplingParams
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import GuidedDecodingParams, RequestOutputKind
+
+from mllm.models.inference_backend import LLMInferenceBackend, PolicyOutput
+from mllm.utils.short_id_gen import generate_short_id
+
+
+class VLLMAsyncBackend(LLMInferenceBackend):
+    def __init__(
+        self,
+        model_name: str,
+        tokenizer: AutoTokenizer,
+        adapter_paths: dict[str, str],
+        engine_init_kwargs: dict = {},
+        sampling_params: dict = {},
+    ):
+        self.model_name = model_name
+        self.adapter_paths = adapter_paths or {}
+        self.current_adapter = None
+        self.vllm_adapter_ids = {
+            adapter_id: generate_short_id() for adapter_id in adapter_paths.keys()
+        }
+        ea = dict(model=model_name, **engine_init_kwargs)
+        ea["enable_lora"] = True
+        ea["max_loras"] = len(self.vllm_adapter_ids)
+        ea["enable_sleep_mode"] = True
+        self.engine = AsyncLLMEngine.from_engine_args(AsyncEngineArgs(**ea))
+
+        self.sampling_params = sampling_params
+
+    def prepare_adapter(
+        self, adapter_id: Optional[str], weights_got_updated: bool
+    ) -> None:
+        self.current_adapter = adapter_id
+        if weights_got_updated:
+            self.vllm_adapter_ids[adapter_id] = generate_short_id()
+        self.current_lora_request = LoRARequest(
+            adapter_id,
+            self.vllm_adapter_ids[adapter_id],
+            self.adapter_paths[adapter_id],
+        )
+
+    async def toggle_training_mode(self) -> None:
+        await self.engine.sleep(level=1)
+
+    async def toggle_eval_mode(self) -> None:
+        await self.engine.wake_up()
+
+    def shutdown(self) -> None:
+        # No explicit close call; engine stops when process exits.
+        pass
+
+    async def generate(
+        self, prompt_text: str, regex: Optional[str] = None
+    ) -> PolicyOutput:
+        # Build SamplingParams correctly
+
+        guided = GuidedDecodingParams(regex=regex) if regex else None
+        sp = SamplingParams(
+            **self.sampling_params,
+            guided_decoding=guided,
+            output_kind=RequestOutputKind.FINAL_ONLY,
+        )
+
+        request_id = f"req-{asyncio.get_running_loop().time()}"
+        result_generator = self.engine.generate(
+            prompt_text,
+            sp,  # SamplingParams(...)
+            request_id,
+            lora_request=self.current_lora_request,
+        )
+
+        async for out in result_generator:  # with FINAL_ONLY this runs once
+            res = out
+
+        raw_text = res.outputs[0].text
+
+        content = raw_text
+        reasoning_content = None
+
+        if regex:
+            # Strict split: require \n<think>...</think>\n\n before final content
+            m = re.match(
+                r"^\n<think>\n([\s\S]*?)</think>\n\n(.*)$", raw_text, flags=re.DOTALL
+            )
+            if m:
+                reasoning_content = m.group(1)
+                content = m.group(2)
+
+        return PolicyOutput(content=content, reasoning_content=reasoning_content)

src_code_for_reproducibility/models/inference_backend_vllm_local_server.py

@@ -0,0 +1,160 @@
+import json
+import os
+import subprocess
+import time
+
+import httpx
+import requests
+
+from mllm.models.inference_backend import LLMInferenceBackend
+
+
+class HttpVLLMBackend(LLMInferenceBackend):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.port = kwargs.get("port", 8000)
+        self.host = kwargs.get("host", "0.0.0.0")
+        self.proc = None
+        self.base_url = f"http://{self.host}:{self.port}"
+        # vLLM memory safety knobs
+        self.gpu_mem_util = kwargs.get("gpu_memory_utilization", 0.9)
+        self.max_model_len = kwargs.get("max_model_len", None)
+        self.max_num_seqs = kwargs.get("max_num_seqs", None)
+        self.max_batched_tokens = kwargs.get("max_num_batched_tokens", None)
+        self.dtype = kwargs.get("dtype", "bfloat16")
+        self.trust_remote_code = kwargs.get("trust_remote_code", False)
+        # LoRA strategy: "preload" (CLI) or "runtime" (endpoints) depending on your vLLM build
+        self.lora_mode = kwargs.get(
+            "lora_mode", "preload"
+        )  # "runtime" supported in newer builds
+        self.runtime_lora_enabled = self.lora_mode == "runtime"
+
+        # If preloading: build CLI args (adapter name -> path)
+        self._preload_lora_args = []
+        if self.adapter_paths and self.lora_mode == "preload":
+            # vLLM supports multiple LoRA modules via CLI in recent versions
+            # Example flag shapes can vary; adapt as needed for your version:
+            # --lora-modules adapter_id=path
+            for aid, pth in self.adapter_paths.items():
+                self._preload_lora_args += ["--lora-modules", f"{aid}={pth}"]
+
+    def launch(self):
+        # Build vLLM serve command
+        cmd = [
+            "python3",
+            "-m",
+            "vllm.entrypoints.openai.api_server",
+            "--model",
+            self.model_name,
+            "--host",
+            self.host,
+            "--port",
+            str(self.port),
+            "--dtype",
+            self.dtype,
+            "--gpu-memory-utilization",
+            str(self.gpu_mem_util),
+        ]
+        if self.trust_remote_code:
+            cmd += ["--trust-remote-code"]
+        if self.max_model_len:
+            cmd += ["--max-model-len", str(self.max_model_len)]
+        if self.max_num_seqs:
+            cmd += ["--max-num-seqs", str(self.max_num_seqs)]
+        if self.max_batched_tokens:
+            cmd += ["--max-num-batched-tokens", str(self.max_batched_tokens)]
+        cmd += self._preload_lora_args
+
+        self.proc = subprocess.Popen(
+            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
+        )
+        self._wait_ready()
+
+    def _wait_ready(self, timeout=120):
+        url = f"{self.base_url}/v1/models"
+        t0 = time.time()
+        while time.time() - t0 < timeout:
+            try:
+                r = requests.get(url, timeout=2)
+                if r.status_code == 200:
+                    return
+            except Exception:
+                pass
+            time.sleep(1)
+        raise RuntimeError("vLLM server did not become ready in time")
+
+    def is_ready(self) -> bool:
+        try:
+            return (
+                requests.get(f"{self.base_url}/v1/models", timeout=2).status_code == 200
+            )
+        except Exception:
+            return False
+
+    def prepare_adapter(self, adapter_id: str) -> None:
+        if not adapter_id or not self.runtime_lora_enabled:
+            return
+        # Newer vLLM builds expose runtime LoRA endpoints. If yours differs,
+        # adjust the path/body here and keep the interface stable.
+        try:
+            requests.post(
+                f"{self.base_url}/v1/load_lora_adapter",
+                json={
+                    "adapter_name": adapter_id,
+                    "adapter_path": self.adapter_paths[adapter_id],
+                },
+                timeout=10,
+            ).raise_for_status()
+        except Exception as e:
+            # If already loaded or endpoint not present, swallow or log
+            pass
+
+    async def generate(
+        self, prompt_text: str, sampling_params: dict, adapter_id: str | None
+    ) -> str:
+        # Map your sampling params to OpenAI schema
+        body = {
+            "model": self.model_name,
+            "messages": [{"role": "user", "content": prompt_text}],
+            "temperature": sampling_params.get("temperature", 1.0),
+            "top_p": sampling_params.get("top_p", 1.0),
+            "max_tokens": sampling_params.get("max_new_tokens", 128),
+        }
+        # Optional knobs:
+        if sampling_params.get("top_k", -1) and sampling_params["top_k"] > 0:
+            # vLLM accepts top_k via extra params; put under "extra_body"
+            body.setdefault("extra_body", {})["top_k"] = sampling_params["top_k"]
+        if sampling_params.get("min_new_tokens", None) is not None:
+            body.setdefault("extra_body", {})["min_tokens"] = sampling_params[
+                "min_new_tokens"
+            ]
+        if sampling_params.get("frequency_penalty", None) is not None:
+            body["frequency_penalty"] = sampling_params["frequency_penalty"]
+
+        # Select LoRA adapter
+        if adapter_id:
+            if self.runtime_lora_enabled:
+                body.setdefault("extra_body", {})["lora_adapter"] = adapter_id
+            else:
+                # when preloaded via CLI, most builds select by name via "adapter_name"/"lora_adapter"
+                body.setdefault("extra_body", {})["lora_adapter"] = adapter_id
+
+        url = f"{self.base_url}/v1/chat/completions"
+        timeout = httpx.Timeout(3600.0, connect=3600.0)
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            resp = await client.post(url, json=body)
+            resp.raise_for_status()
+            data = resp.json()
+            return data["choices"][0]["message"]["content"]
+
+    def toggle_training_mode(self) -> None:
+        # vLLM doesn’t expose an explicit KV “release” toggle via API.
+        # Strategy: keep inference server idle during training, or run training in a separate process.
+        pass
+
+    def toggle_eval_mode(self) -> None:
+        pass
+
+    def shutdown(self) -> None:
+        if self.proc:
+            self.proc.terminate()