Add files using upload-large-folder tool

seed_1/Qwen/Qwen2.5-7B-Instruct/adapters/fixed_ad_align_adapter/adapter_config.json

@@ -0,0 +1,46 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": null,
+  "base_model_name_or_path": "Qwen/Qwen2.5-7B-Instruct",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 64,
+  "lora_bias": false,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "peft_version": "0.18.1",
+  "qalora_group_size": 16,
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "down_proj",
+    "up_proj",
+    "gate_proj",
+    "o_proj",
+    "k_proj",
+    "q_proj",
+    "v_proj"
+  ],
+  "target_parameters": null,
+  "task_type": "CAUSAL_LM",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false
+}

src_code_for_reproducibility/chat_utils/apply_template.py

@@ -0,0 +1,89 @@
+"""
+File: mllm/chat_utils/apply_template.py
+Summary: Applies tokenizer-specific chat templates and stitches chat token IDs.
+"""
+
+import torch
+
+from mllm.chat_utils.chat_turn import ChatTurn
+from mllm.chat_utils.template_specific import (
+    custom_gemma3_template,
+    custom_llama3_template,
+    custom_qwen2_template,
+    custom_qwen3_template,
+    gemma3_assistant_postfix,
+    qwen2_assistant_postfix,
+    qwen3_assistant_postfix,
+)
+
+
+def get_custom_chat_template(tokenizer) -> str:
+    """
+    Get the chat template for the tokenizer.
+    """
+    if "qwen2" in tokenizer.name_or_path.lower():
+        return custom_qwen2_template
+    elif "llama" in tokenizer.name_or_path.lower():
+        return custom_llama3_template
+    elif "qwen3" in tokenizer.name_or_path.lower():
+        return custom_qwen3_template
+    elif "gemma" in tokenizer.name_or_path.lower():
+        return custom_gemma3_template
+    else:
+        raise ValueError(f"Tokenizer {tokenizer.name_or_path} not supported")
+
+
+def get_custom_assistant_postfix(tokenizer) -> torch.Tensor:
+    """
+    Get the custom assistant postfix for the tokenizer.
+    """
+    if "qwen2" in tokenizer.name_or_path.lower():
+        return qwen2_assistant_postfix
+    elif "qwen3" in tokenizer.name_or_path.lower():
+        return qwen3_assistant_postfix
+    elif "gemma" in tokenizer.name_or_path.lower():
+        return gemma3_assistant_postfix
+    return torch.tensor([], dtype=torch.long)
+
+
+def tokenize_chats(chats: list[ChatTurn], tokenizer, enable_thinking) -> None:
+    """
+    Set the chat_template_token_ids for each chat turn.
+    We rely on tokenizer-side templates because engine-provided cached tokens are not exposed yet.
+    """
+    custom_template = get_custom_chat_template(tokenizer)
+    custom_assistant_postfix: torch.Tensor = get_custom_assistant_postfix(tokenizer)
+    for i, chat in enumerate(chats):
+        if chat.chat_template_token_ids is None:
+            if chat.role == "user":
+                next_chat = chats[i + 1] if i + 1 < len(chats) else None
+                add_generation_prompt = True
+                if next_chat and next_chat.role == "user":
+                    add_generation_prompt = False
+                encoded_chat = tokenizer.apply_chat_template(
+                    [chat],
+                    return_tensors="pt",
+                    chat_template=custom_template,
+                    add_generation_prompt=add_generation_prompt,
+                    add_system_prompt=True if i == 0 else False,
+                    enable_thinking=enable_thinking,
+                ).flatten()
+                previous_chat = chats[i - 1] if i > 0 else None
+                if previous_chat and previous_chat.role == "assistant":
+                    encoded_chat = torch.cat([custom_assistant_postfix, encoded_chat])
+            elif chat.role == "assistant":
+                encoded_chat = chat.out_token_ids
+            chat.chat_template_token_ids = encoded_chat
+
+
+def chat_turns_to_token_ids(
+    chats: list[ChatTurn], tokenizer, enable_thinking
+) -> list[int]:
+    """
+    Tokenize the chat turns and set the chat_template_token_ids for each chat turn.
+    """
+    tokenize_chats(chats=chats, tokenizer=tokenizer, enable_thinking=enable_thinking)
+    token_ids = []
+    for chat in chats:
+        token_ids.append(chat.chat_template_token_ids)
+    return torch.cat(token_ids)

src_code_for_reproducibility/chat_utils/chat_turn.py

@@ -0,0 +1,32 @@
+"""
+File: mllm/chat_utils/chat_turn.py
+Summary: Defines the ChatTurn schema plus helpers for serialization and validation.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, List, Literal, Optional, Tuple
+
+import jsonschema
+import torch
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+AgentId = str
+
+
+class ChatTurn(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)  # needed for torch tensors
+
+    role: str = Field(pattern="^(user|assistant)$")
+    agent_id: AgentId  # ID of the agent with which the chat occured
+    content: str
+    reasoning_content: str | None = None
+    chat_template_token_ids: torch.LongTensor | None = None  # Token ids of chat template format. For example, token ids of "<assistant>{content}</assistant>""
+    out_token_ids: torch.LongTensor | None = (
+        None  # tokens generated from inference engine
+    )
+    log_probs: torch.FloatTensor | None = None
+    is_state_end: bool = False  # indicates whether this chat turn marks the end of a state in the trajectory

src_code_for_reproducibility/chat_utils/template_specific.py

@@ -0,0 +1,114 @@
+"""
+File: mllm/chat_utils/template_specific.py
+Summary: Stores chat template variants and assistant postfix tensors per tokenizer.
+"""
+
+import huggingface_hub
+import torch
+from transformers import AutoTokenizer
+
+custom_llama3_template = """
+{%- if add_system_prompt %}
+    {{- '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nCutting Knowledge Date: December 2023\nToday Date: 26 Jul 2024\n\n<|eot_id|>' }}
+{%- endif %}
+{%- for message in messages %}
+    {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+    {{- '<|start_header_id|>' + 'assistant' + '<|end_header_id|>\n\n' }}
+{%- endif %}
+"""
+
+qwen2_assistant_postfix = (
+    AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
+    .encode("\n", return_tensors="pt")
+    .flatten()
+)
+qwen3_assistant_postfix = (
+    AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
+    .encode("\n", return_tensors="pt")
+    .flatten()
+)
+gemma3_assistant_postfix = (
+    AutoTokenizer.from_pretrained("google/gemma-3-4b-it")
+    .encode("\n", return_tensors="pt")
+    .flatten()
+)
+custom_qwen2_template = """
+{%- if add_system_prompt %}
+    {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if reasoning_content %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}
+"""
+
+custom_qwen3_template = """
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}
+"""
+
+custom_gemma3_template = """
+{%- if add_system_prompt %}
+{{- bos_token -}}
+{%- endif %}
+{%- for message in messages -%}
+{%- if message['role'] == 'assistant' -%}
+{%- set role = 'model' -%}
+{%- else -%}
+{%- set role = message['role'] -%}
+{%- endif -%}
+{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+{{ '<start_of_turn>model\n' }}
+{%- endif -%}
+"""

src_code_for_reproducibility/markov_games/__init__.py

@@ -0,0 +1,4 @@
+"""
+File: mllm/markov_games/__init__.py
+Summary: Makes Markov-game subpackages importable from the top-level namespace.
+"""

src_code_for_reproducibility/markov_games/agent.py

@@ -0,0 +1,72 @@
+"""
+File: mllm/markov_games/agent.py
+Summary: Declares the base Agent interface connecting simulations to policy calls.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Tuple
+
+from numpy.random import default_rng
+
+from mllm.markov_games.rollout_tree import AgentActLog
+
+
+class Agent(ABC):
+    """Abstract policy wrapper that bridges simulations with arbitrary backends."""
+
+    @abstractmethod
+    def __init__(
+        self,
+        seed: int,
+        agent_id: str,
+        agent_name: str,
+        agent_policy: Callable[[list[dict]], str],
+        *args,
+        **kwargs,
+    ):
+        """
+        Initialize the agent state and seed its RNG.
+
+        Subclasses typically store extra handles (tokenizers, inference clients, etc.)
+        but they should always call ``super().__init__`` so sampling remains reproducible.
+        """
+        self.seed = seed
+        self.agent_id = agent_id
+        self.agent_name = agent_name
+        self.policy = policy
+        self.rng = default_rng(self.seed)
+        raise NotImplementedError
+
+    async def act(self, observation) -> Tuple[Any, AgentActLog]:
+        """
+        Produce the next action (and associated chat log) given an environment observation.
+
+        Implementations can iterate with rejection sampling, multi-call deliberation, etc.
+        Returns both the chosen action and an `AgentActLog` describing how it was produced.
+        """
+        raise NotImplementedError
+
+    def get_safe_copy(self):
+        """
+        Return a deep copy whose future calls do not mutate the original agent.
+
+        Needed for branch exploration/reruns with alternative actions.
+        """
+        raise NotImplementedError
+
+    def reset(self):
+        """Reset any internal state between rollouts."""
+        raise NotImplementedError
+
+    def render(self):
+        """Optional human-readable visualization of the agent (CLI/UI)."""
+        raise NotImplementedError
+
+    def close(self):
+        """Release any external resources (network sockets, subprocesses, etc.)."""
+        raise NotImplementedError
+
+    def get_agent_info(self):
+        """Return diagnostic metadata to embed inside rollout logs."""
+        raise NotImplementedError

src_code_for_reproducibility/markov_games/alternative_actions_runner.py

@@ -0,0 +1,146 @@
+"""
+File: mllm/markov_games/alternative_actions_runner.py
+Summary: Generates rollout branches by replaying trajectories with unilateral action changes.
+"""
+
+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,
+):
+    """
+    Roll out a counterfactual branch where ``agent_id`` deviates unilaterally.
+
+    Starting from ``branch_node`` (which already contains the main trajectory),
+    we replay the simulation with the deviating agent's action while freezing
+    all other agents/actions, then continue for ``max_depth`` steps.
+    """
+
+    # 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,
+):
+    """
+    Generate a rollout tree containing the main path plus unilateral deviation branches.
+
+    For each timestep we:
+      1. Cache agent actions without side effects.
+      2. Advance the main trajectory.
+      3. Spawn ``nb_alternative_actions`` asynchronous deviations per agent,
+         each replaying up to ``max_depth`` steps from the cached pre-action state.
+    The resulting branches feed advantage-alignment estimators.
+    """
+
+    tasks = []
+    time_step = 0
+    terminated = False
+    root = RolloutTreeRootNode(id=markov_game.get_id(), crn_id=markov_game.get_crn_id())
+    previous_node = root
+
+    while not terminated:
+        mg_before_action = markov_game.get_safe_copy()
+
+        # Get safe copies for main branch
+        agent_action_safe_copies: dict[
+            AgentId, AgentAndActionSafeCopy
+        ] = await markov_game.get_actions_of_agents_without_side_effects()
+
+        markov_game.set_actions_of_agents_manually(agent_action_safe_copies)
+        terminated = markov_game.take_simulation_step()
+        main_node = RolloutTreeNode(
+            step_log=markov_game.get_step_log(), time_step=time_step
+        )
+        branch_node = RolloutTreeBranchNode(main_child=main_node)
+        previous_node.child = branch_node
+        previous_node = main_node
+
+        # Get alternative branches by generating new unilateral actions
+        for agent_id in markov_game.agent_ids:
+            for _ in range(nb_alternative_actions):
+                # Get safe copies for branches
+                branch_agent_action_safe_copies: dict[
+                    AgentId, AgentAndActionSafeCopy
+                ] = {
+                    agent_id: AgentAndActionSafeCopy(
+                        action=copy.deepcopy(agent_action_safe_copy.action),
+                        action_info=copy.deepcopy(agent_action_safe_copy.action_info),
+                        agent_after_action=agent_action_safe_copy.agent_after_action.get_safe_copy(),
+                    )
+                    for agent_id, agent_action_safe_copy in agent_action_safe_copies.items()
+                }
+                mg_branch: MarkovGame = mg_before_action.get_safe_copy()
+                other_agent_id = [id for id in mg_branch.agent_ids if id != agent_id][0]
+                mg_branch.set_action_and_agent_after_action_manually(
+                    agent_id=other_agent_id,
+                    agent_action_safe_copy=branch_agent_action_safe_copies[
+                        other_agent_id
+                    ],
+                )
+                task = asyncio.create_task(
+                    run_with_unilateral_alt_action(
+                        markov_game=mg_branch,
+                        time_step=time_step,
+                        agent_id=agent_id,
+                        branch_node=branch_node,
+                        max_depth=max_depth,
+                    )
+                )
+                tasks.append(task)
+        time_step += 1
+
+    # wait for all branches to complete
+    await asyncio.gather(*tasks)
+
+    return root

src_code_for_reproducibility/markov_games/group_timesteps.py

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

src_code_for_reproducibility/markov_games/linear_runner.py

@@ -0,0 +1,42 @@
+"""
+File: mllm/markov_games/linear_runner.py
+Summary: Simulates a single unbranched Markov-game rollout and records it.
+"""
+
+import asyncio
+import json
+import os.path
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import RolloutTreeNode, RolloutTreeRootNode
+
+
+async def LinearRunner(
+    markov_game: MarkovGame, output_folder: str
+) -> RolloutTreeRootNode:
+    """
+    Generate a single main-path rollout (no branching) for the provided Markov game.
+
+    Parameters
+    ----------
+    markov_game:
+        Initialized ``MarkovGame`` with agents + simulation ready to step.
+    output_folder:
+        Unused placeholder in the legacy API (kept for compatibility).
+    """
+    time_step = 0
+    terminated = False
+    root = RolloutTreeRootNode(
+        id=markov_game.get_id(),
+        crn_id=markov_game.get_crn_id(),
+        agent_ids=markov_game.get_agent_ids(),
+    )
+    previous_node = root
+    while not terminated:
+        terminated, step_log = await markov_game.step()
+        current_node = RolloutTreeNode(step_log=step_log, time_step=time_step)
+        previous_node.child = current_node
+        previous_node = current_node
+        time_step += 1
+
+    return root

src_code_for_reproducibility/markov_games/markov_game.py

@@ -0,0 +1,217 @@
+"""
+File: mllm/markov_games/markov_game.py
+Summary: Defines the MarkovGame base class plus shared simulation interfaces.
+"""
+
+import asyncio
+import copy
+import json
+import os
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, Tuple
+
+from transformers.models.idefics2 import Idefics2Config
+
+from mllm.markov_games.agent import Agent
+from mllm.markov_games.rollout_tree import AgentActLog, StepLog
+from mllm.markov_games.simulation import Simulation
+
+AgentId = str
+
+
+@dataclass
+class AgentAndActionSafeCopy:
+    """Snapshot of an agent, its action, and metadata used for branch replay."""
+
+    action: Any
+    action_info: AgentActLog
+    agent_after_action: type[Agent]
+
+
+class MarkovGame(object):
+    def __init__(
+        self,
+        id: int,
+        agents: dict[AgentId, type[Agent]],
+        simulation: type[Simulation],
+        crn_id: int,
+    ):
+        """
+        Initialize the Markov game wrapper.
+
+        Parameters
+        ----------
+        id:
+            Unique rollout identifier (logged into rollout trees).
+        agents:
+            Mapping of agent_id -> Agent instance.
+        simulation:
+            Environment implementing the ``Simulation`` interface (IPD, TAS, etc.).
+        crn_id:
+            Identifier for the common random number stream used by this rollout.
+        """
+        self.agents = agents
+        self.agent_ids = self.agents.keys()
+        self.simulation = simulation
+        self.simulation_step_log = None
+        self.agent_step_logs = {agent_id: None for agent_id in self.agent_ids}
+        self.actions = {}
+        self.id = id
+        self.crn_id = crn_id
+
+    def get_id(self) -> str:
+        return self.id
+
+    def get_crn_id(self) -> int:
+        return self.crn_id
+
+    def get_agent_ids(self) -> List[AgentId]:
+        return list(self.agent_ids)
+
+    async def get_action_of_agent_without_side_effects(
+        self, agent_id: AgentId
+    ) -> Tuple[Any, AgentActLog]:
+        """
+        Safe function to get an action of an agent without modifying the agent or the simulation.
+        """
+        agent = self.agents[agent_id]
+        agent_before_action = agent.get_safe_copy()
+        obs = self.simulation.get_obs_agent(agent_id)
+        action, action_info = await agent.act(observation=obs)
+        self.agents[agent_id] = agent_before_action
+        agent_after_action = agent.get_safe_copy()
+        return AgentAndActionSafeCopy(action, action_info, agent_after_action)
+
+    async def get_actions_of_agents_without_side_effects(
+        self,
+    ) -> dict[AgentId, AgentAndActionSafeCopy]:
+        """
+        Safe function to get an action of an agent without modifying the agent or the simulation.
+        """
+        tasks = []
+        for agent_id in self.agent_ids:
+            task = asyncio.create_task(
+                self.get_action_of_agent_without_side_effects(agent_id)
+            )
+            tasks.append(task)
+        agent_and_action_safe_copies: list[
+            AgentAndActionSafeCopy
+        ] = await asyncio.gather(*tasks)
+        return {
+            agent_id: agent_and_action_safe_copy
+            for agent_id, agent_and_action_safe_copy in zip(
+                self.agent_ids, agent_and_action_safe_copies
+            )
+        }
+
+    def set_action_and_agent_after_action_manually(
+        self,
+        agent_id: AgentId,
+        agent_action_safe_copy: AgentAndActionSafeCopy,
+    ):
+        """
+        Set the action and the agent after action manually.
+        """
+        self.actions[agent_id] = agent_action_safe_copy.action
+        self.agent_step_logs[agent_id] = agent_action_safe_copy.action_info
+        self.agents[agent_id] = agent_action_safe_copy.agent_after_action
+
+    def set_actions_of_agents_manually(
+        self, actions: dict[AgentId, AgentAndActionSafeCopy]
+    ):
+        """
+        Set the actions of agents manually.
+        """
+        for agent_id, agent_action_safe_copy in actions.items():
+            self.set_action_and_agent_after_action_manually(
+                agent_id, agent_action_safe_copy
+            )
+
+    async def set_action_of_agent(self, agent_id: AgentId):
+        """
+        Query a single agent for its next action and store the result locally.
+        """
+        agent = self.agents[agent_id]
+        obs = self.simulation.get_obs_agent(agent_id)
+        action, action_info = await agent.act(observation=obs)
+        self.actions[agent_id] = action
+        self.agent_step_logs[agent_id] = action_info
+
+    async def set_actions(self):
+        """
+        Query every agent concurrently and populate the cached actions/logs.
+        """
+        # background_tasks = set()
+        tasks = []
+        for agent_id in self.agent_ids:
+            task = asyncio.create_task(self.set_action_of_agent(agent_id))
+            tasks.append(task)
+        await asyncio.gather(*tasks)
+
+    def take_simulation_step(self):
+        """
+        Advance the simulation by one step using the cached actions.
+        """
+        terminated, self.simulation_step_log = self.simulation.step(self.actions)
+        return terminated
+
+    def get_step_log(self) -> StepLog:
+        """
+        Package the most recent simulation step and agent logs into a StepLog.
+        """
+        if self.simulation_step_log is None:
+            raise RuntimeError(
+                "Simulation step log is empty; call take_simulation_step() first."
+            )
+        missing_logs = [
+            agent_id for agent_id, log in self.agent_step_logs.items() if log is None
+        ]
+        if missing_logs:
+            raise RuntimeError(
+                f"Agent action logs missing for: {', '.join(missing_logs)}. "
+                "Ensure set_actions() ran before requesting the step log."
+            )
+        step_log = StepLog(
+            simulation_step_log=self.simulation_step_log,
+            action_logs=self.agent_step_logs,
+        )
+        return step_log
+
+    async def step(self) -> Tuple[bool, StepLog]:
+        """
+        Convenience step that collects actions, advances the simulation, and returns the log.
+        """
+        await self.set_actions()
+        terminated = self.take_simulation_step()
+        step_log = self.get_step_log()
+        return terminated, step_log
+
+    def get_safe_copy(self):
+        """
+        Create a shallow copy of the game with deep-copied agents/simulation for branching.
+        """
+
+        new_markov_game = copy.copy(self)
+        new_simulation = self.simulation.get_safe_copy()
+        new_agents = {
+            agent_id: agent.get_safe_copy() for agent_id, agent in self.agents.items()
+        }
+
+        # Reassign copied components
+        new_markov_game.simulation = new_simulation
+        new_markov_game.agents = new_agents
+
+        # IMPORTANT: ensure agent_ids references the new agents dict, not the original
+        new_markov_game.agent_ids = new_markov_game.agents.keys()
+
+        # Deep-copy step data to avoid correlation
+        new_markov_game.simulation_step_log = copy.deepcopy(self.simulation_step_log)
+        new_markov_game.actions = copy.deepcopy(self.actions)
+        # Rebuild logs to align exactly with new agent ids
+        old_agent_step_logs = copy.deepcopy(self.agent_step_logs)
+        new_markov_game.agent_step_logs = {
+            agent_id: old_agent_step_logs.get(agent_id)
+            for agent_id in new_markov_game.agent_ids
+        }
+
+        return new_markov_game

src_code_for_reproducibility/markov_games/mg_utils.py

@@ -0,0 +1,97 @@
+"""
+File: mllm/markov_games/mg_utils.py
+Summary: Holds miscellaneous helpers shared across Markov-game modules.
+"""
+
+import asyncio
+import copy
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from mllm.markov_games.ipd.ipd_agent import IPDAgent
+from mllm.markov_games.ipd.Ipd_hard_coded_agents import (
+    AlwaysCooperateIPDAgent,
+    AlwaysDefectIPDAgent,
+)
+from mllm.markov_games.ipd.ipd_simulation import IPD
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.negotiation.dond_agent import DealNoDealAgent
+from mllm.markov_games.negotiation.dond_simulation import DealNoDealSimulation
+from mllm.markov_games.negotiation.nego_hard_coded_policies import (
+    HardCodedNegoGreedyPolicy,
+    HardCodedNegoWelfareMaximizingPolicy,
+)
+from mllm.markov_games.negotiation.no_press_nego_agent import NoPressAgent
+from mllm.markov_games.negotiation.no_press_nego_simulation import NoPressSimulation
+from mllm.markov_games.negotiation.tas_rps_agent import TrustAndSplitRPSAgent
+from mllm.markov_games.negotiation.tas_rps_simulation import TrustAndSplitRPSSimulation
+from mllm.markov_games.rollout_tree import (
+    AgentActLog,
+    RolloutTreeBranchNode,
+    RolloutTreeNode,
+    RolloutTreeRootNode,
+    StepLog,
+)
+from mllm.markov_games.simulation import SimulationStepLog
+
+AgentId = str
+
+
+@dataclass
+class AgentConfig:
+    """Configuration blob describing one agent in a Markov game spec."""
+
+    agent_id: str
+    agent_name: str
+    agent_class_name: str
+    policy_id: str
+    init_kwargs: dict
+
+
+@dataclass
+class MarkovGameConfig:
+    """Top-level config that ties together simulation settings and agent configs."""
+
+    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]]
+):
+    """
+    Materialize Agents and the Simulation described by ``config`` and return a MarkovGame.
+
+    `policies` is a mapping of policy_id -> callable retrieved from the hosting trainer.
+    """
+    agents = {}
+    agent_names = []
+    for agent_config in config.agent_configs:
+        agent_id = agent_config.agent_id
+        agent_name = agent_config.agent_name
+        agent_class = eval(agent_config.agent_class_name)
+        agent = agent_class(
+            seed=config.seed,
+            agent_id=agent_id,
+            agent_name=agent_name,
+            policy=policies[agent_config.policy_id],
+            **agent_config.init_kwargs,
+        )
+        agents[agent_id] = agent
+        agent_names.append(agent_name)
+    simulation = eval(config.simulation_class_name)(
+        seed=config.seed,
+        agent_ids=list(agents.keys()),
+        agent_names=agent_names,
+        **config.simulation_init_args,
+    )
+    markov_game = MarkovGame(
+        id=config.id,
+        crn_id=config.seed,
+        agents=agents,
+        simulation=simulation,
+    )
+    return markov_game

src_code_for_reproducibility/markov_games/rollout_tree.py

@@ -0,0 +1,95 @@
+"""
+File: mllm/markov_games/rollout_tree.py
+Summary: Defines rollout tree data structures and serialization helpers.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, List, Literal, Optional, Tuple
+
+import jsonschema
+from pydantic import BaseModel, Field, model_validator
+
+from mllm.chat_utils.chat_turn import ChatTurn
+
+AgentId = str
+
+
+class SimulationStepLog(BaseModel):
+    """Minimal snapshot of environment-side rewards and auxiliary info."""
+
+    rewards: dict[AgentId, float]
+    info: Any = None
+
+
+class AgentActLog(BaseModel):
+    """LLM-side provenance for an action (chat turns + metadata)."""
+
+    chat_turns: list[ChatTurn] | None
+    info: Any = None
+
+    @model_validator(mode="after")
+    def _exactly_one_state_end(self):
+        """
+        This method is used to enforce that for each AgentActLog, there is exactly one ChatTurn which is a state end.
+        """
+        if self.chat_turns != []:
+            n = sum(1 for t in self.chat_turns if t.is_state_end)
+            if n != 1:
+                raise ValueError(
+                    f"AgentActLog must have exactly one ChatTurn with is_state_end=True; got {self.chat_turns}."
+                )
+            return self
+        else:
+            return self
+
+
+class StepLog(BaseModel):
+    action_logs: dict[AgentId, AgentActLog]
+    simulation_step_log: SimulationStepLog
+
+
+# BranchType = Literal["unilateral_deviation", "common_deviation"] # might not be necessary
+# class BranchNodeInfo(BaseModel):
+#     branch_id: str
+#     branch_for: AgentId
+#     branch_type: BranchType
+
+
+class RolloutTreeNode(BaseModel):
+    """Single timestep of the main trajectory (or a branch) plus linkage."""
+
+    step_log: StepLog
+    time_step: int
+    child: RolloutTreeNode | RolloutTreeBranchNode | None = None
+
+
+class RolloutTreeBranchNode(BaseModel):
+    """
+    First item of the tuple indicates which agent "called" for an alternative branch.
+    """
+
+    main_child: RolloutTreeNode
+    branches: dict[AgentId, list[RolloutTreeNode]] | None = None
+
+
+class RolloutTreeRootNode(BaseModel):
+    """Entry point for serialized rollouts (main path plus optional branches)."""
+
+    id: int
+    crn_id: int  # ID of the rng used to generate this rollout tree
+    child: RolloutTreeNode | RolloutTreeBranchNode | None = None
+    agent_ids: List[AgentId] = Field(min_length=1)
+
+
+# class RolloutTreeLeafNode(BaseModel):
+#     step_log: StepLog
+#     time_step: int
+
+
+# Necessary for self-referential stuff in pydantic
+RolloutTreeBranchNode.model_rebuild()
+RolloutTreeNode.model_rebuild()

src_code_for_reproducibility/markov_games/run_markov_games.py

@@ -0,0 +1,52 @@
+"""
+File: mllm/markov_games/run_markov_games.py
+Summary: CLI entry point for running configured Markov-game experiments.
+"""
+
+import asyncio
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from torch._C import ClassType
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import RolloutTreeRootNode
+
+
+async def run_markov_games(
+    runner: Callable[[MarkovGame], RolloutTreeRootNode],
+    runner_kwargs: dict,
+    output_folder: str,
+    markov_games: list[MarkovGame],
+) -> list[RolloutTreeRootNode]:
+    """
+    Kick off multiple Markov game rollouts concurrently and return their trees.
+
+    Parameters mirror the Hydra configs (runner callable + kwargs) so callers can
+    choose ``LinearRunner``, ``AlternativeActionsRunner`` or future variants.
+    """
+    runner_kwargs = dict(runner_kwargs)
+    max_parallel_games = runner_kwargs.pop("max_parallel_games", None)
+
+    async def run_game(markov_game: MarkovGame) -> RolloutTreeRootNode:
+        return await runner(
+            markov_game=markov_game,
+            output_folder=output_folder,
+            **runner_kwargs,
+        )
+
+    if max_parallel_games is not None:
+        semaphore = asyncio.Semaphore(max(1, int(max_parallel_games)))
+
+        async def run_game(markov_game: MarkovGame) -> RolloutTreeRootNode:
+            async with semaphore:
+                return await runner(
+                    markov_game=markov_game,
+                    output_folder=output_folder,
+                    **runner_kwargs,
+                )
+
+    tasks = []
+    for mg in markov_games:
+        tasks.append(asyncio.create_task(run_game(mg)))
+    return await asyncio.gather(*tasks)

src_code_for_reproducibility/markov_games/simulation.py

@@ -0,0 +1,94 @@
+"""
+File: mllm/markov_games/simulation.py
+Summary: Core simulation loop utilities and step logging for Markov games.
+"""
+
+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]:
+        """
+        Advance the environment by one logical tick using ``actions``.
+
+        Returns
+        -------
+        terminated: bool
+            Whether the episode has finished.
+        SimulationStepLog
+            Reward/info bundle describing this transition.
+        """
+        raise NotImplementedError
+
+    def get_obs(self):
+        """Return a dict mapping agent_id -> observation for *all* agents."""
+        raise NotImplementedError
+
+    def get_obs_agent(self, agent_id):
+        """Return the observation for a single agent."""
+        raise NotImplementedError
+
+    def get_obs_size(self):
+        """Describe the observation tensor shape (useful for critic heads)."""
+        raise NotImplementedError
+
+    def get_state(self):
+        """Return the privileged simulator state if available."""
+        raise NotImplementedError
+
+    def get_state_size(self):
+        """Describe the state tensor shape."""
+        raise NotImplementedError
+
+    def get_avail_actions(self):
+        """Return the global action mask/tensor if the space is discrete."""
+        raise NotImplementedError
+
+    def get_avail_agent_actions(self, agent_id):
+        """Return the available action mask for a given agent."""
+        raise NotImplementedError
+
+    def get_total_actions(self):
+        """Returns the total number of actions an agent could ever take.
+
+        Implementations currently assume a discrete, one-dimensional action space per agent.
+        """
+        raise NotImplementedError
+
+    def get_safe_copy(self):
+        """
+        Return copy of the simulator that shares no mutable state with the original.
+        """
+        raise NotImplementedError
+
+    def reset(self):
+        """Reset to the initial state and return the starting observations."""
+        raise NotImplementedError
+
+    def render(self):
+        """Optional human-facing visualization."""
+        raise NotImplementedError
+
+    def close(self):
+        """Release any owned resources (files, processes, etc.)."""
+        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,415 @@
+"""
+File: mllm/markov_games/statistics_runner.py
+Summary: Executes multiple rollouts to compute experiment statistics.
+"""
+
+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]:
+    """Yield ``SimulationStepLog`` objects along the main (non-branch) path."""
+    for node in _iterate_main_nodes(root):
+        yield node.step_log.simulation_step_log
+
+
+def stream_rollout_files(iteration_folder: Path) -> Iterator[Path]:
+    """Iterate over every ``*.rt.pkl`` file under an iteration directory."""
+    for p in iteration_folder.rglob("*.rt.pkl"):
+        if p.is_file():
+            yield p
+
+
+def load_root(path: Path) -> RolloutTreeRootNode:
+    """Load and validate a rollout tree from disk."""
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+    return RolloutTreeRootNode.model_validate(data)
+
+
+@dataclass
+class StatRecord:
+    """Convenience container for serialized stat rows."""
+
+    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/__init__.py

@@ -0,0 +1,4 @@
+"""
+File: mllm/models/__init__.py
+Summary: Exports model-layer utilities from the models package.
+"""

src_code_for_reproducibility/models/adapter_training_wrapper.py

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

src_code_for_reproducibility/models/human_policy.py

@@ -0,0 +1,260 @@
+"""
+File: mllm/models/human_policy.py
+Summary: Implements an interactive human-in-the-loop policy for experiments.
+"""
+
+import asyncio
+import os
+import re
+import shutil
+import sys
+from typing import Callable, Dict, List, Optional
+
+from mllm.markov_games.rollout_tree import ChatTurn
+
+try:
+    import rstr  # For generating example strings from regex
+except Exception:  # pragma: no cover
+    rstr = None
+
+
+def _clear_terminal() -> None:
+    """
+    Clear the terminal screen in a cross-platform manner.
+    """
+    if sys.stdout.isatty():
+        os.system("cls" if os.name == "nt" else "clear")
+
+
+def _terminal_width(default: int = 100) -> int:
+    try:
+        return shutil.get_terminal_size().columns
+    except Exception:
+        return default
+
+
+def _horizontal_rule(char: str = "─") -> str:
+    width = max(20, _terminal_width() - 2)
+    return char * width
+
+
+class _Style:
+    # ANSI colors (bright, readable)
+    RESET = "\033[0m"
+    BOLD = "\033[1m"
+    DIM = "\033[2m"
+    # Foreground colors
+    FG_BLUE = "\033[94m"  # user/system headers
+    FG_GREEN = "\033[92m"  # human response header
+    FG_YELLOW = "\033[93m"  # notices
+    FG_RED = "\033[91m"  # errors
+    FG_MAGENTA = "\033[95m"  # regex
+    FG_CYAN = "\033[96m"  # tips
+
+
+def _render_chat(state) -> str:
+    """
+    Render prior messages in a compact, readable terminal format.
+
+    Expected message dict keys: {"role": str, "content": str, ...}
+    """
+    lines: List[str] = []
+    lines.append(_horizontal_rule())
+    lines.append(f"{_Style.FG_BLUE}{_Style.BOLD} Conversation so far {_Style.RESET}")
+    lines.append(_horizontal_rule())
+    for chat in state:
+        role = chat.role
+        content = str(chat.content).strip()
+        # Map roles to display names and colors/emojis
+        if role == "assistant":
+            header = f"{_Style.FG_GREEN}{_Style.BOLD}HUMAN--🧑‍💻{_Style.RESET}"
+        elif role == "user":
+            header = f"{_Style.FG_BLUE}{_Style.BOLD}USER--⚙️{_Style.RESET}"
+        else:
+            header = f"[{_Style.DIM}{role.upper()}{_Style.RESET}]"
+        lines.append(header)
+        # Indent content for readability
+        for line in content.splitlines() or [""]:
+            lines.append(f"  {line}")
+        lines.append("")
+    lines.append(_horizontal_rule())
+    return "\n".join(lines)
+
+
+async def _async_input(prompt_text: str) -> str:
+    """Non-blocking input using a background thread."""
+    return await asyncio.to_thread(input, prompt_text)
+
+
+def _short_regex_example(regex: str, max_len: int = 30) -> Optional[str]:
+    """
+    Try to produce a short example string that matches the regex.
+    We attempt multiple times and pick the first <= max_len.
+    """
+    if rstr is None:
+        return None
+    try:
+        for _ in range(20):
+            candidate = rstr.xeger(regex)
+            if len(candidate) <= max_len:
+                return candidate
+        # Fallback to truncation (may break match, so don't return)
+        return None
+    except Exception:
+        return None
+
+
+def _detect_input_type(regex: str | None) -> tuple[str, str, str]:
+    """
+    Detect what type of input is expected based on the regex pattern.
+    Returns (input_type, start_tag, end_tag)
+    """
+    if regex is None:
+        return "text", "", ""
+
+    if "message_start" in regex and "message_end" in regex:
+        return "message", "<<message_start>>", "<<message_end>>"
+    elif "proposal_start" in regex and "proposal_end" in regex:
+        return "proposal", "<<proposal_start>>", "<<proposal_end>>"
+    else:
+        return "text", "", ""
+
+
+async def human_policy(state, agent_id, regex: str | None = None) -> str:
+    """
+    Async human-in-the-loop policy.
+
+    - Displays prior conversation context in the terminal.
+    - Prompts the user for a response.
+    - If a regex is provided, validates and re-prompts until it matches.
+    - Automatically adds formatting tags based on expected input type.
+
+    Args:
+        prompt: Chat history as a list of {role, content} dicts.
+        regex: Optional fullmatch validation pattern.
+
+    Returns:
+        The user's validated response string.
+    """
+    # Detect input type and formatting
+    input_type, start_tag, end_tag = _detect_input_type(regex)
+
+    while True:
+        _clear_terminal()
+        print(_render_chat(state))
+
+        if regex:
+            example = _short_regex_example(regex, max_len=30)
+            print(
+                f"{_Style.FG_MAGENTA}{_Style.BOLD}Expected format (regex fullmatch):{_Style.RESET}"
+            )
+            print(f"  {_Style.FG_MAGENTA}{regex}{_Style.RESET}")
+            if example:
+                print(
+                    f"{_Style.FG_CYAN}Example (random, <=30 chars):{_Style.RESET} {example}"
+                )
+            print(_horizontal_rule("."))
+
+            # Custom prompt based on input type
+            if input_type == "message":
+                print(
+                    f"{_Style.FG_YELLOW}Type your message content (formatting will be added automatically):{_Style.RESET}"
+                )
+            elif input_type == "proposal":
+                print(
+                    f"{_Style.FG_YELLOW}Type your proposal (number only, formatting will be added automatically):{_Style.RESET}"
+                )
+            else:
+                print(
+                    f"{_Style.FG_YELLOW}Type your response and press Enter.{_Style.RESET}"
+                )
+
+            print(
+                f"{_Style.DIM}Commands: /help to view commands, /refresh to re-render, /quit to abort{_Style.RESET}"
+            )
+        else:
+            print(
+                f"{_Style.FG_YELLOW}Type your response and press Enter.{_Style.RESET} {_Style.DIM}(/help for commands){_Style.RESET}"
+            )
+
+        user_in = (await _async_input("> ")).rstrip("\n")
+
+        # Commands
+        if user_in.strip().lower() in {"/help", "/h"}:
+            print(f"\n{_Style.FG_CYAN}{_Style.BOLD}Available commands:{_Style.RESET}")
+            print(
+                f"  {_Style.FG_CYAN}/help{_Style.RESET} or {_Style.FG_CYAN}/h{_Style.RESET}     Show this help"
+            )
+            print(
+                f"  {_Style.FG_CYAN}/refresh{_Style.RESET} or {_Style.FG_CYAN}/r{_Style.RESET}  Re-render the conversation and prompt"
+            )
+            print(
+                f"  {_Style.FG_CYAN}/quit{_Style.RESET} or {_Style.FG_CYAN}/q{_Style.RESET}     Abort the run (raises KeyboardInterrupt)"
+            )
+            await asyncio.sleep(1.0)
+            continue
+        if user_in.strip().lower() in {"/refresh", "/r"}:
+            continue
+        if user_in.strip().lower() in {"/quit", "/q"}:
+            raise KeyboardInterrupt("Human aborted run from human_policy")
+
+        # Add formatting tags if needed
+        if start_tag and end_tag:
+            formatted_input = f"{start_tag}{user_in}{end_tag}"
+        else:
+            formatted_input = user_in
+
+        if regex is None:
+            return ChatTurn(
+                role="assistant", agent_id=agent_id, content=formatted_input
+            )
+
+        # Validate against regex (fullmatch)
+        try:
+            pattern = re.compile(regex)
+        except re.error as e:
+            # If regex is invalid, fall back to accepting any input
+            print(
+                f"{_Style.FG_RED}Warning:{_Style.RESET} Provided regex is invalid: {e}. Accepting input without validation."
+            )
+            await asyncio.sleep(0.5)
+            return ChatTurn(
+                role="assistant", agent_id=agent_id, content=formatted_input
+            )
+
+        if pattern.fullmatch(formatted_input):
+            return ChatTurn(
+                role="assistant", agent_id=agent_id, content=formatted_input
+            )
+
+        # Show validation error and re-prompt
+        print("")
+        print(
+            f"{_Style.FG_RED}{_Style.BOLD}Input did not match the required format.{_Style.RESET} Please try again."
+        )
+
+        if input_type == "message":
+            print(
+                f"You entered: {_Style.FG_CYAN}{start_tag}{user_in}{end_tag}{_Style.RESET}"
+            )
+            print(f"Just type the message content without tags.")
+        elif input_type == "proposal":
+            print(
+                f"You entered: {_Style.FG_CYAN}{start_tag}{user_in}{end_tag}{_Style.RESET}"
+            )
+            print(f"Just type the number without tags.")
+        else:
+            print(f"Expected (regex):")
+            print(f"  {_Style.FG_MAGENTA}{regex}{_Style.RESET}")
+
+        print(_horizontal_rule("."))
+        print(f"{_Style.FG_YELLOW}Press Enter to retry...{_Style.RESET}")
+        await _async_input("")
+
+
+def get_human_policies() -> Dict[str, Callable[[List[Dict]], str]]:
+    """
+    Expose the human policy in the same map shape used elsewhere.
+    """
+    # Type hint says Callable[[List[Dict]], str] but we intentionally return the async callable.
+    return {"human_policy": human_policy}  # type: ignore[return-value]

src_code_for_reproducibility/models/inference_backend.py

@@ -0,0 +1,44 @@
+"""
+File: mllm/models/inference_backend.py
+Summary: Declares the inference backend interface and shared dataclasses.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+@dataclass
+class LLMInferenceOutput:
+    content: str
+    reasoning_content: str | None = None
+    log_probs: list[float] | None = None
+    out_token_ids: list[int] | None = None
+
+
+class LLMInferenceBackend(ABC):
+    @abstractmethod
+    def __init__(self, **kwargs):
+        ...
+
+    @abstractmethod
+    def prepare_adapter(
+        self, adapter_id: str, weights_got_updated: bool = False
+    ) -> None:
+        """Ensure adapter is ready/loaded for next generation call."""
+
+    @abstractmethod
+    async def generate(self, prompt: list[dict], regex: Optional[str] = None) -> str:
+        ...
+
+    @abstractmethod
+    def toggle_training_mode(self) -> None:
+        ...
+
+    @abstractmethod
+    def toggle_eval_mode(self) -> None:
+        ...
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        ...

src_code_for_reproducibility/models/inference_backend_dummy.py

@@ -0,0 +1,59 @@
+"""
+File: mllm/models/inference_backend_dummy.py
+Summary: Stub inference backend that returns synthetic completions for tests.
+"""
+
+import asyncio
+from typing import Optional
+
+import rstr
+from transformers import AutoTokenizer
+
+from mllm.models.inference_backend import LLMInferenceBackend, LLMInferenceOutput
+from mllm.utils.short_id_gen import generate_short_id
+
+
+class DummyInferenceBackend(LLMInferenceBackend):
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        pass
+
+    def prepare_adapter(
+        self,
+        adapter_id: Optional[str],
+        weights_got_updated: bool,
+        adapter_path: Optional[str] = None,
+    ) -> None:
+        pass
+
+    async def toggle_training_mode(self) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def toggle_eval_mode(self) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    async def generate(
+        self,
+        prompt_text: str,
+        regex: Optional[str] = None,
+        extract_thinking: bool = False,
+    ) -> LLMInferenceOutput:
+        if regex:
+            # Create random string that respects the regex
+            return LLMInferenceOutput(
+                content=rstr.xeger(regex),
+                reasoning_content="I don't think, I am a dummy backend.",
+            )
+        else:
+            return LLMInferenceOutput(
+                content="I am a dummy backend without a regex.",
+                reasoning_content="I don't think, I am a dummy backend.",
+            )

src_code_for_reproducibility/models/inference_backend_vllm.py

@@ -0,0 +1,111 @@
+"""
+File: mllm/models/inference_backend_vllm.py
+Summary: Connects to in-process vLLM instances for batched generation.
+"""
+
+import asyncio
+import re
+from typing import Optional
+
+import torch
+from transformers import AutoTokenizer
+from vllm import AsyncEngineArgs, AsyncLLMEngine, SamplingParams
+from vllm.inputs import TokensPrompt
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import GuidedDecodingParams, RequestOutputKind
+
+from mllm.models.inference_backend import LLMInferenceBackend, LLMInferenceOutput
+from mllm.utils.short_id_gen import generate_short_id
+
+
+class VLLMAsyncBackend(LLMInferenceBackend):
+    def __init__(
+        self,
+        model_name: str,
+        tokenizer: AutoTokenizer,
+        # adapter_paths: dict[str, str],
+        engine_init_kwargs: dict = {},
+        sampling_params: dict = {},
+    ):
+        self.model_name = model_name
+        self.vllm_adapter_ids = {}
+        ea = dict(model=model_name, **engine_init_kwargs)
+        self.engine = AsyncLLMEngine.from_engine_args(AsyncEngineArgs(**ea))
+
+        self.sampling_params = sampling_params
+        self.tokenizer = tokenizer
+
+    def prepare_adapter(
+        self,
+        adapter_id: Optional[str],
+        adapter_path: Optional[str],
+        weights_got_updated: bool,
+    ) -> None:
+        if weights_got_updated:
+            self.vllm_adapter_ids[adapter_id] = generate_short_id()
+        self.current_lora_request = LoRARequest(
+            adapter_id,
+            self.vllm_adapter_ids[adapter_id],
+            adapter_path,
+        )
+
+    async def toggle_training_mode(self) -> None:
+        await self.engine.sleep(level=1)
+
+    async def toggle_eval_mode(self) -> None:
+        await self.engine.wake_up()
+
+    def shutdown(self) -> None:
+        # No explicit close call; engine stops when process exits.
+        pass
+
+    async def generate(
+        self,
+        input_token_ids: list[int],
+        regex: Optional[str] = None,
+        extract_thinking: bool = False,
+    ) -> LLMInferenceOutput:
+        # Build SamplingParams correctly
+        guided = GuidedDecodingParams(regex=regex) if regex else None
+        sp = SamplingParams(
+            **self.sampling_params,
+            guided_decoding=guided,
+            output_kind=RequestOutputKind.FINAL_ONLY,
+        )
+
+        prompt = TokensPrompt(prompt_token_ids=input_token_ids)
+        request_id = f"req-{asyncio.get_running_loop().time()}"
+        result_generator = self.engine.generate(
+            prompt,
+            sp,  # SamplingParams(...)
+            request_id,
+            lora_request=self.current_lora_request,
+        )
+
+        async for out in result_generator:  # with FINAL_ONLY this runs once
+            res = out
+
+        raw_text = res.outputs[0].text
+        out_token_ids = res.outputs[0].token_ids
+        log_probs = [
+            logprob_dict[token_id].logprob
+            for token_id, logprob_dict in zip(out_token_ids, res.outputs[0].logprobs)
+        ]
+        log_probs = torch.tensor(log_probs)
+        out_token_ids = torch.tensor(out_token_ids, dtype=torch.long)
+        content = raw_text
+        reasoning_content = None
+
+        if extract_thinking:
+            m = re.match(
+                r"^\n<think>\n([\s\S]*?)</think>\n\n(.*)$", raw_text, flags=re.DOTALL
+            )
+            if m:
+                reasoning_content = m.group(1)
+                content = m.group(2)
+        return LLMInferenceOutput(
+            content=content,
+            reasoning_content=reasoning_content,
+            log_probs=log_probs,
+            out_token_ids=out_token_ids,
+        )

src_code_for_reproducibility/models/large_language_model_api.py

@@ -0,0 +1,184 @@
+"""
+File: mllm/models/large_language_model_api.py
+Summary: Implements API-based large-language-model inference adapters.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import os
+import random
+import re
+from typing import Any, Callable, Dict, List, Optional, Sequence
+
+import backoff
+from openai import AsyncOpenAI, OpenAIError
+
+from mllm.markov_games.rollout_tree import ChatTurn
+from mllm.models.inference_backend import LLMInferenceOutput
+
+# Static list copied from the public OpenAI docs until a discovery endpoint is exposed.
+reasoning_models = [
+    "gpt-5-nano",
+    "gpt-5-mini",
+    "gpt-5",
+    "o1-mini",
+    "o1",
+    "o1-pro",
+    "o3-mini",
+    "o3",
+    "o3-pro",
+    "o4-mini",
+    "o4",
+    "o4-pro",
+    "openai/gpt-oss-20b",
+    "openai/gpt-oss-120b",
+]
+
+
+class LargeLanguageModelOpenAI:
+    """Tiny async wrapper for OpenAI Chat Completions."""
+
+    def __init__(
+        self,
+        llm_id: str = "",
+        model: str = "gpt-4.1-mini",
+        reasoning_effort: str = "low",
+        add_constraint_msg: bool = True,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout_s: float = 300.0,
+        regex_max_attempts: int = 10,
+        sampling_params: Optional[Dict[str, Any]] = None,
+        init_kwargs: Optional[Dict[str, Any]] = None,
+        output_directory: Optional[str] = None,
+    ) -> None:
+        self.llm_id = llm_id
+        self.model = model
+        key = api_key or os.getenv("OPENAI_API_KEY")
+        if not key:
+            raise RuntimeError(
+                "Set OPENAI_API_KEY as global environment variable or pass api_key."
+            )
+        client_kwargs: Dict[str, Any] = {"api_key": key, "timeout": timeout_s}
+        if base_url:
+            client_kwargs["base_url"] = base_url
+        self.client = AsyncOpenAI(**client_kwargs)
+
+        # Sampling/default request params set at init
+        self.sampling_params = sampling_params
+        self.use_reasoning = model in reasoning_models
+        if self.use_reasoning:
+            self.sampling_params["reasoning"] = {
+                "effort": reasoning_effort,
+                "summary": "detailed",
+            }
+        self.regex_max_attempts = max(1, int(regex_max_attempts))
+        self.add_constraint_msg = add_constraint_msg
+
+    def get_inference_policies(self) -> Dict[str, Callable]:
+        return {
+            self.llm_id: self.get_action,
+        }
+
+    async def prepare_adapter_for_inference(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def toggle_eval_mode(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def toggle_training_mode(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def export_adapters(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def checkpoint_all_adapters(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    def extract_output_from_response(self, resp: Response) -> LLMInferenceOutput:
+        if len(resp.output) > 1:
+            reasoning_content = resp.output[0].content
+            summary = resp.output[0].summary
+            if reasoning_content is not None:
+                reasoning_content = (
+                    f"OpenAI Reasoning Content: {reasoning_content[0].text}"
+                )
+            elif summary != []:
+                reasoning_content = f"OpenAI Reasoning Summary: {summary[0].text}"
+            else:
+                reasoning_content = None
+            content = resp.output[1].content[0].text
+        else:
+            reasoning_content = None
+            content = resp.output[0].content[0].text
+
+        return LLMInferenceOutput(
+            content=content,
+            reasoning_content=reasoning_content,
+        )
+
+    @backoff.on_exception(
+        backoff.expo, Exception, max_time=10**10, max_tries=10**10
+    )
+    async def get_action(
+        self,
+        state: list[ChatTurn],
+        agent_id: str,
+        regex: Optional[str] = None,
+    ) -> LLMInferenceOutput:
+        # Remove any non-role/content keys from the prompt else openai will error.
+        prompt = [{"role": p.role, "content": p.content} for p in state]
+
+        # if self.sleep_between_requests:
+        #     await self.wait_random_time()
+
+        # If regex is required, prime the model and validate client-side
+        if regex:
+            if self.add_constraint_msg:
+                constraint_msg = {
+                    "role": "user",
+                    "content": (
+                        f"Output must match this regex exactly: {regex} \n"
+                        "Return only the matching string, with no quotes or extra text."
+                    ),
+                }
+                prompt = [constraint_msg, *prompt]
+            pattern = re.compile(regex)
+            for _ in range(self.regex_max_attempts):
+                resp = await self.client.responses.create(
+                    model=self.model,
+                    input=prompt,
+                    **self.sampling_params,
+                )
+                policy_output = self.extract_output_from_response(resp)
+                if pattern.fullmatch(policy_output.content):
+                    return policy_output
+                prompt = [
+                    *prompt,
+                    {
+                        "role": "user",
+                        "content": (
+                            f"Invalid response format. Expected format (regex): {regex}\n Please try again and provide ONLY a response that matches this regex."
+                        ),
+                    },
+                ]
+            return policy_output
+
+        # Simple, unconstrained generation
+        resp = await self.client.responses.create(
+            model=self.model,
+            input=prompt,
+            **self.sampling_params,
+        )
+        policy_output = self.extract_output_from_response(resp)
+        return policy_output
+
+    def shutdown(self) -> None:
+        self.client = None

src_code_for_reproducibility/models/large_language_model_gemini_api.py

@@ -0,0 +1,197 @@
+"""
+File: mllm/models/large_language_model_gemini_api.py
+Summary: Implements native Gemini API-based large-language-model inference adapters.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import re
+from typing import Any, Callable, Dict, List, Optional
+
+import backoff
+from google import genai
+from google.genai import types
+
+from mllm.markov_games.rollout_tree import ChatTurn
+from mllm.models.inference_backend import LLMInferenceOutput
+
+
+class LargeLanguageModelGemini:
+    """Tiny async wrapper for the native Gemini API."""
+
+    def __init__(
+        self,
+        llm_id: str = "",
+        model: str = "gemini-3.1-flash-lite-preview",
+        api_key: Optional[str] = None,
+        timeout_s: float = 300.0,
+        regex_max_attempts: int = 10,
+        sampling_params: Optional[Dict[str, Any]] = None,
+        thinking_level: str = "low",
+        include_thoughts: bool = True,
+        init_kwargs: Optional[Dict[str, Any]] = None,
+        output_directory: Optional[str] = None,
+    ) -> None:
+        self.llm_id = llm_id
+        self.model = model
+        self.timeout_s = timeout_s
+        key = api_key or os.getenv("GEMINI_API_KEY")
+        if not key:
+            raise RuntimeError(
+                "Set GEMINI_API_KEY as global environment variable or pass api_key."
+            )
+        self.client = genai.Client(api_key=key)
+        self.sampling_params = sampling_params or {}
+        self.thinking_level = thinking_level
+        self.include_thoughts = include_thoughts
+        self.regex_max_attempts = max(1, int(regex_max_attempts))
+
+    def get_inference_policies(self) -> Dict[str, Callable]:
+        return {
+            self.llm_id: self.get_action,
+        }
+
+    async def prepare_adapter_for_inference(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def toggle_eval_mode(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def toggle_training_mode(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def export_adapters(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    async def checkpoint_all_adapters(self, *args: Any, **kwargs: Any) -> None:
+        await asyncio.sleep(0)
+        pass
+
+    def messages_to_contents(self, messages: List[Dict[str, str]]) -> List[types.Content]:
+        contents: List[types.Content] = []
+        system_chunks: List[str] = []
+
+        for message in messages:
+            role = message["role"]
+            text = message["content"]
+
+            if role == "system":
+                system_chunks.append(text)
+                continue
+
+            gemini_role = "model" if role == "assistant" else "user"
+            contents.append(
+                types.Content(
+                    role=gemini_role,
+                    parts=[types.Part.from_text(text=text)],
+                )
+            )
+
+        if system_chunks:
+            system_text = "\n\n".join(system_chunks)
+            contents.insert(
+                0,
+                types.Content(
+                    role="user",
+                    parts=[
+                        types.Part.from_text(
+                            text=(
+                                "System instruction:\n"
+                                f"{system_text}\n\n"
+                                "Follow the system instruction for the rest of this conversation."
+                            )
+                        )
+                    ],
+                ),
+            )
+
+        return contents
+
+    def build_generate_config(self) -> types.GenerateContentConfig:
+        return types.GenerateContentConfig(
+            thinking_config=types.ThinkingConfig(
+                thinking_level=self.thinking_level,
+                include_thoughts=self.include_thoughts,
+            ),
+            **self.sampling_params,
+        )
+
+    def extract_output_from_response(self, response: Any) -> LLMInferenceOutput:
+        reasoning_parts: List[str] = []
+        content_parts: List[str] = []
+
+        if response.candidates:
+            for part in response.candidates[0].content.parts:
+                text = getattr(part, "text", None)
+                if not text:
+                    continue
+                if getattr(part, "thought", False):
+                    reasoning_parts.append(text)
+                else:
+                    content_parts.append(text)
+
+        content = "\n".join(content_parts) if content_parts else (response.text or "")
+        reasoning_content = "\n".join(reasoning_parts) if reasoning_parts else None
+
+        return LLMInferenceOutput(
+            content=content,
+            reasoning_content=reasoning_content,
+        )
+
+    @backoff.on_exception(
+        backoff.expo, Exception, max_time=10**10, max_tries=10**10
+    )
+    async def get_action(
+        self,
+        state: list[ChatTurn],
+        agent_id: str,
+        regex: Optional[str] = None,
+    ) -> LLMInferenceOutput:
+        prompt = [{"role": p.role, "content": p.content} for p in state]
+
+        if regex:
+            constraint_msg = {
+                "role": "user",
+                "content": (
+                    f"Output must match this regex exactly: {regex} \n"
+                    "Return only the matching string, with no quotes or extra text."
+                ),
+            }
+            prompt = [constraint_msg, *prompt]
+            pattern = re.compile(regex)
+            for _ in range(self.regex_max_attempts):
+                response = await self.client.aio.models.generate_content(
+                    model=self.model,
+                    contents=self.messages_to_contents(prompt),
+                    config=self.build_generate_config(),
+                )
+                policy_output = self.extract_output_from_response(response)
+                if pattern.fullmatch(policy_output.content):
+                    return policy_output
+                prompt = [
+                    *prompt,
+                    {
+                        "role": "user",
+                        "content": (
+                            f"Invalid response format. Expected format (regex): {regex}\n"
+                            "Please try again and provide ONLY a response that matches this regex."
+                        ),
+                    },
+                ]
+            return policy_output
+
+        response = await self.client.aio.models.generate_content(
+            model=self.model,
+            contents=self.messages_to_contents(prompt),
+            config=self.build_generate_config(),
+        )
+        return self.extract_output_from_response(response)
+
+    def shutdown(self) -> None:
+        self.client = None

src_code_for_reproducibility/models/large_language_model_local.py

@@ -0,0 +1,361 @@
+"""
+File: mllm/models/large_language_model_local.py
+Summary: Provides a local large language model wrapper over inference backends.
+"""
+
+import logging
+import os
+import re
+import sys
+import uuid
+from collections.abc import Callable
+from copy import deepcopy
+from datetime import datetime
+from typing import Literal
+
+import httpx
+import requests
+import torch
+import torch.nn as nn
+from torch.optim import SGD, Adam, AdamW, RMSprop
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from mllm.chat_utils.apply_template import chat_turns_to_token_ids
+from mllm.markov_games.rollout_tree import ChatTurn
+from mllm.models.adapter_training_wrapper import AdapterWrapper
+from mllm.models.inference_backend import LLMInferenceOutput
+from mllm.models.inference_backend_dummy import DummyInferenceBackend
+from mllm.models.inference_backend_vllm import VLLMAsyncBackend
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+AdapterID = str
+PolicyID = str
+
+
+class LeanLocalLLM:
+    """
+    Wrapper that manages local HuggingFace models, adapters, and inference backends.
+    """
+
+    def __init__(
+        self,
+        llm_id: str = "base_llm",
+        model_name: str = "Qwen/Qwen3-4B-Instruct-2507",
+        device: str = "cuda",
+        hf_kwargs: dict = {},
+        adapter_configs: dict = {},
+        output_directory: str = "./models/",
+        inference_backend: Literal["vllm", "dummy"] = "vllm",
+        inference_backend_sampling_params: dict = {},
+        inference_backend_init_kwargs: dict = {},
+        initial_adapter_paths: dict[str, str] | None = None,
+        initial_buffer_paths: list[str] | None = None,
+        enable_thinking: bool = None,
+        regex_max_attempts: int = -1,
+        max_thinking_characters: int = 0,
+    ):
+        self.inference_backend_name = inference_backend
+        self.output_directory = output_directory
+        self.llm_id = llm_id
+        self.device = torch.device(device) if device else torch.device("cuda")
+        self.model_name = model_name
+        self.adapter_configs = adapter_configs
+        self.adapter_ids = list(adapter_configs.keys())
+        self.enable_thinking = enable_thinking
+        self.regex_max_attempts = regex_max_attempts
+        self.initial_buffer_paths = initial_buffer_paths
+        self.max_thinking_characters = max_thinking_characters
+        self.regex_retries_count = 0
+
+        # Optional user-specified initial adapter weight locations (local or HF Hub)
+        # Format: {adapter_id: path_or_repo_id}
+        self.initial_adapter_paths: dict[str, str] | None = initial_adapter_paths
+
+        # Path management / imports
+        self.save_path = str(os.path.join(output_directory, model_name, "adapters"))
+        self.adapter_paths = {
+            adapter_id: os.path.join(self.save_path, adapter_id)
+            for adapter_id in self.adapter_ids
+        }
+        checkpoints_dir = os.path.join(self.output_directory, "checkpoints")
+        self.past_agent_adapter_paths = {}
+        if os.path.isdir(checkpoints_dir):
+            for dirname in os.listdir(checkpoints_dir):
+                dirpath = os.path.join(checkpoints_dir, dirname)
+                if os.path.isdir(dirpath):
+                    self.past_agent_adapter_paths[f"{dirname}_buffer"] = os.path.join(
+                        dirpath, "agent_adapter"
+                    )
+            logger.info(
+                f"Loaded {len(self.past_agent_adapter_paths)} past agent adapters from checkpoints directory."
+            )
+        if self.initial_buffer_paths is not None:
+            previous_count = len(self.past_agent_adapter_paths)
+            for path in self.initial_buffer_paths:
+                if os.path.isdir(path):
+                    for dirname in os.listdir(path):
+                        dirpath = os.path.join(path, dirname)
+                        if os.path.isdir(dirpath):
+                            self.past_agent_adapter_paths[
+                                f"{dirname}_buffer"
+                            ] = os.path.join(dirpath, "agent_adapter")
+                else:
+                    logger.warning(
+                        f"Initial buffer path {path} does not exist or is not a directory."
+                    )
+            logger.info(
+                f"Loaded {len(self.past_agent_adapter_paths) - previous_count} past agent adapters from user-specified initial buffer paths."
+            )
+        self.past_agent_adapter_ids = list(self.past_agent_adapter_paths.keys())
+
+        # ID management for tracking adapter versions
+        self.adapter_train_ids = {
+            adapter_id: self.short_id_generator() for adapter_id in self.adapter_ids
+        }
+        # Initialize tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        # Setup padding token to be same as EOS token
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        self.weights_got_updated: dict[AdapterID, bool] = {
+            adapter_id: False for adapter_id in self.adapter_ids
+        }
+        self.weights_got_updated.update(
+            {adapter_id: False for adapter_id in self.past_agent_adapter_ids}
+        )
+        self.current_lora_request = None
+        self.currently_loaded_adapter_id = None
+
+        # ---------------------------------------------------------
+        # Init HF model, peft adapters
+        # ---------------------------------------------------------
+        self.shared_hf_llm = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path=model_name,
+            **hf_kwargs,
+        )
+        self.hf_adapters = {}
+        self.optimizers = {}
+        for adapter_id in self.adapter_ids:
+            # Prefer output-folder path if it exists; else fall back to user-specified initial path if provided
+            output_path = os.path.join(self.save_path, adapter_id)
+            chosen_path: str | None = None
+            if os.path.isdir(output_path) and os.listdir(output_path):
+                chosen_path = output_path
+                logger.info(
+                    f"Initializing adapter '{adapter_id}': using existing weights from output folder '{chosen_path}'."
+                )
+            elif (
+                self.initial_adapter_paths and adapter_id in self.initial_adapter_paths
+            ):
+                chosen_path = self.initial_adapter_paths[adapter_id]
+                logger.info(
+                    f"Initializing adapter '{adapter_id}': using provided initial path '{chosen_path}'."
+                )
+            else:
+                logger.info(
+                    f"Initializing adapter '{adapter_id}': no initial weights provided or found; starting from scratch."
+                )
+            hf_adapter = AdapterWrapper(
+                shared_llm=self.shared_hf_llm,
+                adapter_id=adapter_id,
+                lora_config=adapter_configs[adapter_id],
+                path=chosen_path,
+            ).to(device)
+            self.hf_adapters[adapter_id] = hf_adapter
+        # Persist current state of all adapters (ensures remote loads are cached to disk)
+        self.export_adapters()
+
+        # ---------------------------------------------------------
+        # Init inference inference_backend
+        # ---------------------------------------------------------
+
+        if inference_backend == "vllm":
+            self.inference_backend = VLLMAsyncBackend(
+                model_name=self.model_name,
+                # adapter_paths=self.adapter_paths,
+                tokenizer=self.tokenizer,
+                engine_init_kwargs=inference_backend_init_kwargs,
+                sampling_params=inference_backend_sampling_params,
+            )
+        elif inference_backend == "dummy":
+            self.inference_backend = DummyInferenceBackend()
+        else:
+            raise ValueError(f"Unknown inference_backend: {inference_backend}")
+
+    def reset_regex_retries_count(self) -> None:
+        self.regex_retries_count = 0
+
+    def get_inference_policies(self) -> dict[PolicyID, Callable]:
+        """
+        Build async policy callables keyed by adapter id for inference-only usage.
+        """
+        policies = {}
+        for adapter_id in self.adapter_ids:
+            # define policy func
+            async def policy(
+                state: list[ChatTurn],
+                agent_id: str,
+                regex: str | None = None,
+                _adapter_id=adapter_id,
+            ):
+                self.prepare_adapter_for_inference(adapter_id=_adapter_id)
+                response = await self.get_action(state, agent_id, regex)
+                return response
+
+            policies[self.llm_id + "/" + adapter_id] = policy
+
+        for adapter_id in self.past_agent_adapter_ids:
+            # define policy func
+            async def policy(
+                state: list[ChatTurn],
+                agent_id: str,
+                regex: str | None = None,
+                _adapter_id=adapter_id,
+            ):
+                self.prepare_adapter_for_inference(adapter_id=_adapter_id)
+                response = await self.get_action(state, agent_id, regex)
+                return response
+
+            policies[self.llm_id + "/" + adapter_id] = policy
+        return policies
+
+    def get_adapter_modules(self) -> dict[PolicyID, nn.Module]:
+        """
+        Returns wrappers over the adapters which allows them be
+        interfaced like regular PyTorch models.
+        AdapterWrapper lives in adapter_wrapper.py; the huggingface modules already wrap
+        parameters here, so we surface them directly until an extra shim is required.
+        """
+        trainable_objects = {an: self.hf_adapters[an] for an in self.adapter_ids}
+        return trainable_objects
+
+    async def toggle_training_mode(self) -> None:
+        for adn in self.adapter_ids:
+            self.adapter_train_ids[adn] = self.short_id_generator()
+        await self.inference_backend.toggle_training_mode()
+
+    async def toggle_eval_mode(self) -> None:
+        await self.inference_backend.toggle_eval_mode()
+
+    def prepare_adapter_for_inference(self, adapter_id: AdapterID) -> None:
+        self.inference_backend.prepare_adapter(
+            adapter_id,
+            adapter_path=self.adapter_paths.get(
+                adapter_id, self.past_agent_adapter_paths.get(adapter_id, None)
+            ),
+            weights_got_updated=self.weights_got_updated[adapter_id],
+        )
+        self.currently_loaded_adapter_id = adapter_id
+        self.weights_got_updated[adapter_id] = False
+
+    # def _make_prompt_text(self, prompt: list[dict]) -> str:
+    #     if self.enable_thinking is not None:
+    #         prompt_text = self.tokenizer.apply_chat_template(
+    #             prompt,
+    #             tokenize=False,
+    #             add_generation_prompt=True,
+    #             enable_thinking=self.enable_thinking,
+    #         )
+    #     else:
+    #         prompt_text = self.tokenizer.apply_chat_template(
+    #             prompt,
+    #             tokenize=False,
+    #             add_generation_prompt=True,
+    #         )
+
+    #     return prompt_text
+
+    async def get_action(
+        self, state: list[ChatTurn], agent_id: str, regex: str | None = None
+    ) -> ChatTurn:
+        current_regex = regex if self.regex_max_attempts == -1 else None
+        pattern = re.compile(regex) if regex else None
+        nb_attempts = 0
+        state = state[:]
+        while True:
+            context_token_ids = chat_turns_to_token_ids(
+                chats=state,
+                tokenizer=self.tokenizer,
+                enable_thinking=self.enable_thinking,
+            )
+            policy_output = await self.inference_backend.generate(
+                input_token_ids=context_token_ids.tolist(),
+                extract_thinking=(self.max_thinking_characters > 0),
+                regex=current_regex,
+            )
+            if (
+                pattern is None
+                or (nb_attempts >= self.regex_max_attempts)
+                or (pattern.fullmatch(policy_output.content))
+            ):
+                return ChatTurn(
+                    agent_id=agent_id,
+                    role="assistant",
+                    content=policy_output.content,
+                    reasoning_content=policy_output.reasoning_content,
+                    out_token_ids=policy_output.out_token_ids,
+                    log_probs=policy_output.log_probs,
+                    is_state_end=False,
+                )
+            else:
+                self.regex_retries_count += 1
+                nb_attempts += 1
+                logger.warning(
+                    f"Response {policy_output.content} did not match regex: {regex}, retry {nb_attempts}/{self.regex_max_attempts}"
+                )
+                if nb_attempts == self.regex_max_attempts:
+                    current_regex = regex
+                # regex_prompt = ChatTurn(
+                #     role="user",
+                #     content=f"Invalid response format. Expected format (regex): {current_regex}\n Please try again and provide ONLY a response that matches this regex.",
+                #     reasoning_content=None,
+                #     log_probs=None,
+                #     out_token_ids=None,
+                #     is_state_end=False,
+                # )
+                # state.append(regex_prompt)
+
+    def export_adapters(self) -> None:
+        """
+        Any peft wrapper, by default, saves all adapters, not just the one currently loaded.
+        """
+
+        # New version of the adapters available
+        for adapter_id in self.adapter_ids:
+            self.weights_got_updated[adapter_id] = True
+        for adapter_id in self.past_agent_adapter_ids:
+            self.weights_got_updated[adapter_id] = True
+
+        adapter_id = self.adapter_ids[0]
+        self.hf_adapters[adapter_id].save_pretrained(self.save_path)
+
+    def checkpoint_all_adapters(self, checkpoint_indicator: str) -> None:
+        """
+        Checkpoints all adapters to the configured output directory.
+        """
+        adapter_id = self.adapter_ids[0]
+        output_dir = os.path.join(self.output_directory, "checkpoints")
+        os.makedirs(output_dir, exist_ok=True)
+        date_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        agent_adapter_dir = f"{adapter_id}-{checkpoint_indicator}-{date_str}"
+        export_path = os.path.join(output_dir, agent_adapter_dir)
+        for adapter_id in self.adapter_ids:
+            if "agent" in adapter_id:
+                self.past_agent_adapter_paths[
+                    f"{agent_adapter_dir}_buffer"
+                ] = os.path.join(export_path, adapter_id)
+                self.past_agent_adapter_ids.append(f"{agent_adapter_dir}_buffer")
+                self.weights_got_updated[f"{agent_adapter_dir}_buffer"] = False
+                self.hf_adapters[adapter_id].save_pretrained(export_path)
+
+    def short_id_generator(self) -> str:
+        """
+        Generates a short unique ID for tracking adapter versions.
+
+        Returns:
+            int: An 8-digit integer ID.
+        """
+        return str(uuid.uuid4().int)[:8]

src_code_for_reproducibility/models/scalar_critic.py

@@ -0,0 +1,59 @@
+"""
+File: mllm/models/scalar_critic.py
+Summary: Defines a scalar critic network and helper utilities.
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from peft import LoraConfig, get_peft_model
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from mllm.models.adapter_training_wrapper import AdapterWrapper
+
+
+class ScalarCritic(nn.Module):
+    """
+    A causal-LM critic_adapter + a scalar value head:
+        V_φ(s) = wᵀ h_last + b
+    Only LoRA adapters (inside critic_adapter) and the value head are trainable.
+    """
+
+    def __init__(self, critic_adapter: AdapterWrapper):
+        super().__init__()
+        self.critic_adapter = critic_adapter
+        hidden_size = self.critic_adapter.shared_llm.config.hidden_size
+        self.value_head = nn.Linear(hidden_size, 1).to(
+            dtype=critic_adapter.dtype, device=critic_adapter.device
+        )
+
+    def forward(self, input_ids, attention_mask=None, **kwargs):
+        # AdapterWrapper activates its own adapter internally
+        outputs = self.critic_adapter(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            **kwargs,
+        )
+        h_last = outputs.hidden_states[-1]  # (B, S, H)
+        values = self.value_head(h_last).squeeze(-1)  # (B, S)
+        return values
+
+    def parameters(self, recurse: bool = True):
+        """Iterator over *trainable* parameters for this critic."""
+        # 1) LoRA params for *this* adapter
+        for p in self.critic_adapter.parameters():
+            yield p
+        # 2) scalar head
+        yield from self.value_head.parameters()
+
+    def gradient_checkpointing_enable(self, *args, **kwargs):
+        self.critic_adapter.gradient_checkpointing_enable(*args, **kwargs)
+
+    @property
+    def dtype(self):
+        return self.critic_adapter.dtype
+
+    @property
+    def device(self):
+        return self.critic_adapter.device

src_code_for_reproducibility/training/__init__.py

@@ -0,0 +1,4 @@
+"""
+File: mllm/training/__init__.py
+Summary: Exposes training submodules through the package namespace.
+"""

src_code_for_reproducibility/training/annealing_methods.py

@@ -0,0 +1,20 @@
+"""
+File: mllm/training/annealing_methods.py
+Summary: Implements annealing schedules used across training loops.
+"""
+
+import numpy as np
+
+
+def sigmoid_annealing(step: int, temperature: float) -> float:
+    """
+    Smoothly ramp a scalar from 0 → 1 using a temperature-controlled sigmoid.
+
+    Args:
+        step: Current training step or iteration.
+        temperature: Controls how sharp the transition is; larger values flatten the curve.
+
+    Returns:
+        Float in [-1, 1] that can be rescaled for annealing schedules.
+    """
+    return 2 / (1 + np.exp(-step / temperature)) - 1

src_code_for_reproducibility/training/credit_methods.py

@@ -0,0 +1,307 @@
+"""
+File: mllm/training/credit_methods.py
+Summary: Holds credit-assignment routines for reinforcement learning updates.
+"""
+
+import torch
+
+
+def whiten_advantages(advantages: torch.Tensor) -> torch.Tensor:
+    """
+    Normalize a vector of advantages to zero mean / unit variance (global).
+
+    Useful for variance reduction before computing gradients.
+    """
+    whitened_advantages = (advantages - torch.mean(advantages)) / (
+        torch.std(advantages) + 1e-9
+    )
+    return whitened_advantages
+
+
+def whiten_advantages_time_step_wise(
+    advantages: torch.Tensor,  # (B, T)
+) -> torch.Tensor:
+    """
+    Whiten advantages independently per timestep (column-wise mean/std).
+
+    Helps when rollout lengths differ or certain positions have very different scales.
+    """
+    assert advantages.dim() == 2, "Wrong dimensions."
+    whitened_advantages_time_step_wise = (
+        advantages - advantages.mean(dim=0, keepdim=True)
+    ) / (advantages.std(dim=0, keepdim=True) + 1e-9)
+    return whitened_advantages_time_step_wise
+
+
+def get_discounted_state_visitation_credits(
+    credits: torch.Tensor, discount_factor: float  # (B, T)
+) -> torch.Tensor:
+    """
+    Apply geometric discounting to credits so earlier visits count less.
+
+    Equivalent to per-timestep multiplication by ``gamma^t``.
+    """
+    return credits * (
+        discount_factor ** torch.arange(credits.shape[1], device=credits.device)
+    )
+
+
+def get_discounted_returns(
+    rewards: torch.Tensor,  # (B, T)
+    discount_factor: float,
+) -> torch.Tensor:
+    """
+    Computes Monte Carlo discounted returns for a sequence of rewards.
+
+    Args:
+        rewards (torch.Tensor): Array of rewards for each timestep.
+
+    Returns:
+        torch.Tensor: Array of discounted returns.
+    """
+    assert rewards.dim() == 2, "Wrong dimensions."
+    B, T = rewards.shape
+    discounted_returns = torch.zeros_like(rewards)
+    accumulator = torch.zeros(B, device=rewards.device, dtype=rewards.dtype)
+    for t in reversed(range(T)):
+        accumulator = rewards[:, t] + discount_factor * accumulator
+        discounted_returns[:, t] = accumulator
+    return discounted_returns
+
+
+def get_rloo_credits(credits: torch.Tensor):  # (B, S)
+    """Compute leave-one-out baselines for a batch of credits."""
+    assert credits.dim() == 2, "Wrong dimensions."
+    rloo_baselines = torch.zeros_like(credits)
+    n = credits.shape[0]
+    if n == 1:
+        return credits, rloo_baselines
+    rloo_baselines = (torch.sum(credits, dim=0, keepdim=True) - credits) / (n - 1)
+    rloo_credits = credits - rloo_baselines
+    return rloo_credits, rloo_baselines
+
+
+def get_generalized_advantage_estimates(
+    rewards: torch.Tensor,  # (B, T)
+    value_estimates: torch.Tensor,  # (B, T+1)
+    discount_factor: float,
+    lambda_coef: float,
+) -> torch.Tensor:
+    """
+    Compute Generalized Advantage Estimates (GAE).
+
+    See https://arxiv.org/pdf/1506.02438 for derivation.
+    """
+    assert rewards.dim() == value_estimates.dim() == 2, "Wrong dimensions."
+
+    assert (
+        rewards.shape[0] == value_estimates.shape[0]
+    ), f"Got shapes {rewards.shape} and {value_estimates.shape} of rewards and value estimates."
+    assert (
+        rewards.shape[1] == value_estimates.shape[1] - 1
+    ), f"Got shapes {rewards.shape} and {value_estimates.shape} of rewards and value estimates."
+
+    T = rewards.shape[1]
+    tds = rewards + discount_factor * value_estimates[:, 1:] - value_estimates[:, :-1]
+    gaes = torch.zeros_like(tds)
+    acc = 0.0
+    for t in reversed(range(T)):
+        acc = tds[:, t] + lambda_coef * discount_factor * acc
+        gaes[:, t] = acc
+    return gaes
+
+
+def get_advantage_alignment_weights(
+    advantages: torch.Tensor,  # (B, T)
+    exclude_k_equals_t: bool,
+    gamma: float,
+    discount_t: bool,
+) -> torch.Tensor:
+    """
+    The advantage alignment credit is calculated as
+
+    \[
+        A^*(s_t, a_t, b_t) = A^1(s_t, a_t, b_t) + \beta \cdot
+        \left( \sum_{k < t} \gamma^{t-k} A^1(s_k, a_k, b_k) \right)
+        A^2(s_t, a_t, b_t)
+    \]
+
+    Here, the weights are defined as \( \beta \cdot
+        \left( \sum_{k < t} \gamma^{t-k} A^1(s_k, a_k, b_k) \)
+    """
+    T = advantages.shape[1]
+    discounted_advantages = advantages * (
+        gamma * torch.ones((1, T), device=advantages.device)
+    ) ** (-torch.arange(0, T, 1, device=advantages.device))
+    if exclude_k_equals_t:
+        sub = torch.eye(T, device=advantages.device)
+    else:
+        sub = torch.zeros((T, T), device=advantages.device)
+    # Identity is for \( k < t \), remove for \( k \leq t \)
+    ad_align_weights = discounted_advantages @ (
+        torch.triu(torch.ones((T, T), device=advantages.device)) - sub
+    )
+    t_discounts = (gamma * torch.ones((1, T), device=advantages.device)) ** (
+        torch.arange(0, T, 1, device=advantages.device)
+    )
+    ad_align_weights = t_discounts * ad_align_weights
+    if discount_t:
+        time_discounted_advantages = advantages * (
+            gamma * torch.ones((1, T), device=advantages.device)
+        ) ** (torch.arange(0, T, 1, device=advantages.device))
+        ad_align_weights = ad_align_weights - advantages + time_discounted_advantages
+    return ad_align_weights
+
+
+def get_advantage_alignment_credits(
+    a1: torch.Tensor,  # (B, S)
+    a1_alternative: torch.Tensor,  # (B, S, A)
+    a2: torch.Tensor,  # (B, S)
+    exclude_k_equals_t: bool,
+    beta: float,
+    gamma: float = 1.0,
+    use_old_ad_align: bool = False,
+    use_sign: bool = False,
+    clipping: float | None = None,
+    use_time_regularization: bool = False,
+    force_coop_first_step: bool = False,
+    use_variance_regularization: bool = False,
+    rloo_branch: bool = False,
+    reuse_baseline: bool = False,
+    mean_normalize_ad_align: bool = False,
+    whiten_adalign_advantages: bool = False,
+    whiten_adalign_advantages_time_step_wise: bool = False,
+    discount_t: bool = False,
+) -> torch.Tensor:
+    """
+    Calculate the advantage alignment credits with vectorization, as described in https://arxiv.org/abs/2406.14662.
+
+    Recall that the advantage opponent shaping term of the AdAlign policy gradient is:
+    \[
+        \beta \mathbb{E}_{\substack{
+        \tau \sim \text{Pr}_{\mu}^{\pi^1, \pi^2} \\
+        a_t' \sim \pi^1(\cdot \mid s_t)
+        }}
+        \left[\sum_{t=0}^\infty  \gamma^{t}\left( \sum_{k\leq t} A^1(s_k,a^{\prime}_k,b_k) \right) A^{2}(s_t,a_t, b_t)\nabla_{\theta^1}\text{log } \pi^1(a_t|s_t) \right]
+    \]
+
+    This method computes the following:
+    \[
+        Credit(s_t, a_t, b_t) = \gamma^t \left[ A^1(s_t, a_t, b_t) + \beta \left( \sum_{k\leq t} A^1(s_k,a^{\prime}_k,b_k) \right) A^{2}(s_t,a_t, b_t) \right]
+    \]
+
+    Args:
+        a1: Advantages of the main trajectories for the current agent.
+        a1_alternative: Advantages of the alternative trajectories for the current agent.
+        a2: Advantages of the main trajectories for the other agent.
+        discount_factor: Discount factor for the advantage alignment.
+        beta: Beta parameter for the advantage alignment.
+        gamma: Gamma parameter for the advantage alignment.
+        use_sign_in_ad_align: Whether to use sign in the advantage alignment.
+
+    Returns:
+        torch.Tensor: The advantage alignment credits.
+    """
+
+    assert a1.dim() == a2.dim() == 2, "Advantages must be of shape (B, S)"
+    if a1_alternative is not None:
+        assert (
+            a1_alternative.dim() == 3
+        ), "Alternative advantages must be of shape (B, S, A)"
+        B, T, A = a1_alternative.shape
+    else:
+        B, T = a1.shape
+    assert a1.shape == a2.shape, "Not the same shape"
+
+    sub_tensors = {}
+
+    if use_old_ad_align:
+        ad_align_weights = get_advantage_alignment_weights(
+            advantages=a1,
+            exclude_k_equals_t=exclude_k_equals_t,
+            gamma=gamma,
+            discount_t=discount_t,
+        )
+        sub_tensors["ad_align_weights_prev"] = ad_align_weights
+        if exclude_k_equals_t:
+            ad_align_weights = gamma * ad_align_weights
+    else:
+        assert a1_alternative is not None, "Alternative advantages must be provided"
+        if rloo_branch:
+            a1_alternative = torch.cat([a1.unsqueeze(2), a1_alternative], dim=2)
+            a1_alternative = a1_alternative.mean(dim=2)
+            a1, baseline = get_rloo_credits(a1)
+            if reuse_baseline:
+                a1_alternative = a1_alternative - baseline
+            else:
+                a1_alternative, _ = get_rloo_credits(a1_alternative)
+        assert a1.shape == a1_alternative.shape, "Not the same shape"
+        ad_align_weights = get_advantage_alignment_weights(
+            advantages=a1_alternative,
+            exclude_k_equals_t=exclude_k_equals_t,
+            gamma=gamma,
+        )
+        sub_tensors["ad_align_weights"] = ad_align_weights
+
+    # Use sign
+    if use_sign:
+        assert beta == 1.0, "beta should be 1.0 when using sign"
+        positive_signs = ad_align_weights > 0
+        negative_signs = ad_align_weights < 0
+        ad_align_weights[positive_signs] = 1
+        ad_align_weights[negative_signs] = -1
+        sub_tensors["ad_align_weights_sign"] = ad_align_weights
+        # (rest are 0)
+
+    ###################
+    # Process weights
+    ###################
+
+    # Use clipping
+    if clipping not in [0.0, None]:
+        upper_mask = ad_align_weights > 1
+        lower_mask = ad_align_weights < -1
+
+        ad_align_weights = torch.clip(
+            ad_align_weights,
+            -clipping,
+            clipping,
+        )
+        clipping_ratio = (
+            torch.sum(upper_mask) + torch.sum(lower_mask)
+        ) / upper_mask.size
+        sub_tensors["clipped_ad_align_weights"] = ad_align_weights
+
+    # 1/1+t Regularization
+    if use_time_regularization:
+        t_values = torch.arange(1, T + 1).to(ad_align_weights.device)
+        ad_align_weights = ad_align_weights / t_values
+        sub_tensors["time_regularized_ad_align_weights"] = ad_align_weights
+
+    # Use coop on t=0
+    if force_coop_first_step:
+        ad_align_weights[:, 0] = 1
+        sub_tensors["coop_first_step_ad_align_weights"] = ad_align_weights
+
+    ####################################
+    # Compose elements together
+    ####################################
+
+    opp_shaping_terms = beta * ad_align_weights * a2
+    sub_tensors["ad_align_opp_shaping_terms"] = opp_shaping_terms
+
+    credits = a1 + opp_shaping_terms
+    if mean_normalize_ad_align:
+        credits = credits - credits.mean(dim=0)
+        sub_tensors["mean_normalized_ad_align_credits"] = credits
+    if whiten_adalign_advantages:
+        credits = (credits - credits.mean()) / (credits.std() + 1e-9)
+        sub_tensors["whitened_ad_align_credits"] = credits
+    if whiten_adalign_advantages_time_step_wise:
+        credits = (credits - credits.mean(dim=0, keepdim=True)) / (
+            credits.std(dim=0, keepdim=True) + 1e-9
+        )
+        sub_tensors["whitened_ad_align_credits_time_step_wise"] = credits
+    sub_tensors["final_ad_align_credits"] = credits
+
+    return credits, sub_tensors

src_code_for_reproducibility/training/tally_metrics.py

@@ -0,0 +1,64 @@
+"""
+File: mllm/training/tally_metrics.py
+Summary: Transforms tally files into aggregated metric summaries.
+"""
+
+import os
+from numbers import Number
+from typing import Union
+
+import wandb
+
+
+class Tally:
+    """
+    Minimal scalar-first tally.
+    - Keys are strings.
+    - First add stores a scalar; subsequent adds upgrade to a list of scalars.
+    """
+
+    def __init__(self):
+        self.stats = {}
+
+    def reset(self):
+        """Reset all recorded metrics back to an empty dictionary."""
+        self.stats = {}
+
+    def _coerce_scalar(self, value: Union[int, float]) -> Union[int, float]:
+        """Ensure ``value`` is a plain Python scalar (detach tensors, etc.)."""
+        if hasattr(value, "item") and callable(getattr(value, "item")):
+            try:
+                value = value.item()
+            except Exception:
+                pass
+        if isinstance(value, Number):
+            return value
+        raise AssertionError("Metric must be a scalar number")
+
+    def add_metric(self, path: str, metric: Union[int, float]):
+        """Accumulate a metric under ``path`` (scalar on first add, list thereafter)."""
+        metric = float(metric)
+        assert isinstance(path, str), "Path must be a string."
+        assert isinstance(metric, float), "Metric must be a scalar number."
+
+        scalar = self._coerce_scalar(metric)
+        existing = self.stats.get(path)
+        if existing is None:
+            self.stats[path] = scalar
+        elif isinstance(existing, list):
+            existing.append(scalar)
+        else:
+            self.stats[path] = [existing, scalar]
+
+    def save(self, identifier: str, folder: str):
+        """Persist the tally as a pickle file under ``folder``."""
+        os.makedirs(name=folder, exist_ok=True)
+        try:
+            import pickle
+
+            pkl_path = os.path.join(folder, f"{identifier}.tally.pkl")
+            payload = self.stats
+            with open(pkl_path, "wb") as f:
+                pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
+        except Exception:
+            pass

src_code_for_reproducibility/training/tally_rollout.py

@@ -0,0 +1,116 @@
+"""
+File: mllm/training/tally_rollout.py
+Summary: Serializes rollout data into tallies for downstream processing.
+"""
+
+import json
+import os
+from copy import deepcopy
+from typing import Union
+
+import numpy as np
+import pandas as pd
+import torch
+from transformers import AutoTokenizer
+
+
+class RolloutTallyItem:
+    def __init__(
+        self,
+        crn_ids: list[str],
+        rollout_ids: list[str],
+        agent_ids: list[str],
+        metric_matrix: torch.Tensor,
+    ):
+        """Lightweight data container that keeps rollout-aligned metric matrices."""
+        if isinstance(crn_ids, torch.Tensor):
+            crn_ids = crn_ids.detach().cpu().numpy()
+        if isinstance(rollout_ids, torch.Tensor):
+            rollout_ids = rollout_ids.detach().cpu().numpy()
+        if isinstance(agent_ids, torch.Tensor):
+            agent_ids = agent_ids.detach().cpu().numpy()
+        self.crn_ids = crn_ids
+        self.rollout_ids = rollout_ids
+        self.agent_ids = agent_ids
+        metric_matrix = metric_matrix.detach().cpu()
+        assert (
+            0 < metric_matrix.ndim <= 2
+        ), "Metric matrix must have less than or equal to 2 dimensions"
+        if metric_matrix.ndim == 1:
+            metric_matrix = metric_matrix.reshape(1, -1)
+        # Convert to float32 if tensor is in BFloat16 format (not supported by numpy)
+        if metric_matrix.dtype == torch.bfloat16:
+            metric_matrix = metric_matrix.float()
+        self.metric_matrix = metric_matrix.numpy()
+
+
+class RolloutTally:
+    """
+    Tally is a utility class for collecting and storing training metrics.
+    It supports adding metrics at specified paths and saving them to disk.
+    """
+
+    def __init__(self):
+        """
+        Initializes the RolloutTally object.
+
+        Args:
+            tokenizer (AutoTokenizer): Tokenizer for converting token IDs to strings.
+            max_context_length (int, optional): Maximum context length for contextualized metrics. Defaults to 30.
+        """
+        # Array-preserving structure (leaf lists hold numpy arrays / scalars)
+        self.metrics = {}
+        # Global ordered list of sample identifiers (crn_id, rollout_id) added in the order samples are processed
+
+    def reset(self):
+        """Reset the tally to an empty dict."""
+        self.metrics = {}
+
+    def get_from_nested_dict(self, dictio: dict, path: str):
+        """Retrieve a nested entry, creating intermediate dicts as needed."""
+        assert isinstance(path, list), "Path must be list."
+        for sp in path[:-1]:
+            dictio = dictio.setdefault(sp, {})
+        return dictio.get(path[-1], None)
+
+    def set_at_path(self, dictio: dict, path: str, value):
+        """Store ``value`` at ``path``; helper used by ``add_metric``."""
+        for sp in path[:-1]:
+            dictio = dictio.setdefault(sp, {})
+        dictio[path[-1]] = value
+
+    def add_metric(self, path: list[str], rollout_tally_item: RolloutTallyItem):
+        """
+        Adds a metric to the base tally at the specified path.
+
+        Args:
+            path (list): List of keys representing the path in the base tally.
+            rollout_tally_item (RolloutTallyItem): The rollout tally item to add.
+        """
+        rollout_tally_item = deepcopy(rollout_tally_item)
+
+        # Update array-preserving tally
+        array_list = self.get_from_nested_dict(dictio=self.metrics, path=path)
+        if array_list is None:
+            self.set_at_path(dictio=self.metrics, path=path, value=[rollout_tally_item])
+        else:
+            array_list.append(rollout_tally_item)
+
+    def save(self, identifier: str, folder: str):
+        """Persist the tally as a pickle (metrics only) under ``folder``."""
+        os.makedirs(name=folder, exist_ok=True)
+
+        from datetime import datetime
+
+        now = datetime.now()
+
+        # Pickle only (fastest, exact structure with numpy/scalars at leaves)
+        try:
+            import pickle
+
+            pkl_path = os.path.join(folder, f"{identifier}.rt_tally.pkl")
+            payload = {"metrics": self.metrics}
+            with open(pkl_path, "wb") as f:
+                pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
+        except Exception:
+            pass

src_code_for_reproducibility/training/tally_tokenwise.py

@@ -0,0 +1,278 @@
+"""
+File: mllm/training/tally_tokenwise.py
+Summary: Converts token-level tallies into per-token statistics.
+"""
+
+import json
+import os
+from typing import Any, Dict, List, Tuple, Union
+
+import numpy as np
+import pandas as pd
+import torch
+from transformers import AutoTokenizer
+
+
+class ContextualizedTokenwiseTally:
+    """
+    Collect, store, and save token-level metrics per rollout.
+
+    - One DataFrame per rollout_id in `paths`
+    - Index = timestep (int)
+    - Columns are added incrementally via `add_contexts()` and `add_data()`
+    - Cells may contain scalars, strings, or lists (dtype=object)
+    """
+
+    def __init__(
+        self,
+        tokenizer: AutoTokenizer,
+        paths: List[str],
+        max_context_length: int = 30,
+    ):
+        """
+        Args:
+            tokenizer: HuggingFace tokenizer used to convert tids -> tokens
+            paths: rollout identifiers (parallel to batch dimension)
+            max_context_length: truncate context token lists to this length
+        """
+        self.tokenizer = tokenizer
+        self.paths = paths
+        self.max_context_length = max_context_length
+        self.tally: Dict[str, pd.DataFrame] = {path: pd.DataFrame() for path in paths}
+
+        # set later by setters
+        self.contexts: torch.Tensor | None = None
+        self.action_mask: torch.Tensor | None = None
+        self.range: Tuple[int, int] | None = None
+
+    # --------- Utilities ---------
+
+    def tids_to_str(self, tids: List[int]) -> List[str]:
+        """Convert a list of token IDs to a list of token strings."""
+        return self.tokenizer.convert_ids_to_tokens(tids)
+
+    def _ensure_ready(self):
+        """Validate that action mask and range are configured prior to writes."""
+        assert self.action_mask is not None, "call set_action_mask(mask) first"
+        assert self.range is not None, "call set_range((start, end)) first"
+
+    @staticmethod
+    def _sanitize_filename(name: Any) -> str:
+        """Make a safe filename from any rollout_id."""
+        s = str(name)
+        bad = {os.sep, " ", ":", "|", "<", ">", '"', "'"}
+        if os.altsep is not None:
+            bad.add(os.altsep)
+        for ch in bad:
+            s = s.replace(ch, "_")
+        return s
+
+    @staticmethod
+    def _pad_left(seq: List[Any], length: int, pad_val: Any = "") -> List[Any]:
+        """Left-pad a sequence to `length` with `pad_val`."""
+        if len(seq) >= length:
+            return seq[-length:]
+        return [pad_val] * (length - len(seq)) + list(seq)
+
+    # --------- Setters ---------
+
+    def set_action_mask(self, action_mask: torch.Tensor):
+        """Register the (B, S) mask indicating which tokens correspond to actions."""
+        self.action_mask = action_mask
+
+    def set_range(self, range: Tuple[int, int]):
+        """Record which subset of ``paths`` the current mini-batch corresponds to."""
+        self.range = range
+
+    # --------- Column builders ---------
+
+    def add_contexts(self, contexts: torch.Tensor):
+        """
+        Add a single 'context' column (list[str]) for valid steps.
+
+        Expects `contexts` with shape (B, S): token id at each timestep.
+        For each valid timestep t, we use the last N tokens up to and including t:
+            window = contexts[i, max(0, t - N + 1) : t + 1]
+        The list is left-padded with "" to always be length N.
+        """
+        self._ensure_ready()
+
+        current_paths = self.paths[self.range[0] : self.range[1]]
+        B, S = contexts.shape
+        N = self.max_context_length
+
+        # to CPU ints once
+        contexts_cpu = contexts.detach().to("cpu")
+
+        for i in range(B):
+            rollout_id = current_paths[i]
+            df = self.tally.get(rollout_id, pd.DataFrame())
+
+            valid_idx = torch.nonzero(
+                self.action_mask[i].bool(), as_tuple=False
+            ).squeeze(-1)
+            if valid_idx.numel() == 0:
+                self.tally[rollout_id] = df
+                continue
+
+            idx_list = valid_idx.tolist()
+
+            # ensure index contains valid steps
+            if df.empty:
+                df = pd.DataFrame(index=idx_list)
+            else:
+                new_index = sorted(set(df.index.tolist()) | set(idx_list))
+                if list(df.index) != new_index:
+                    df = df.reindex(new_index)
+
+            # build context windows
+            ctx_token_lists = []
+            for t in idx_list:
+                start = max(0, t - N + 1)
+                window_ids = contexts_cpu[i, start : t + 1].tolist()
+                window_toks = self.tids_to_str([int(x) for x in window_ids])
+                if len(window_toks) < N:
+                    window_toks = [""] * (N - len(window_toks)) + window_toks
+                else:
+                    window_toks = window_toks[-N:]
+                ctx_token_lists.append(window_toks)
+
+            # single 'context' column
+            if "context" not in df.columns:
+                df["context"] = pd.Series(index=df.index, dtype=object)
+            df.loc[idx_list, "context"] = pd.Series(
+                ctx_token_lists, index=idx_list, dtype=object
+            )
+
+            self.tally[rollout_id] = df
+
+    def add_data(
+        self,
+        metric_id: str,
+        metrics: torch.Tensor,
+        to_tids: bool = False,
+    ):
+        """
+        Add a metric column for valid steps.
+
+        Args:
+            metric_id: column name
+            metrics: shape (B, S) for scalars/ids or (B, S, K) for top-k vectors
+            to_tids: if True, treat ints/lists of ints as tids and convert to tokens
+        """
+        self._ensure_ready()
+        current_paths = self.paths[self.range[0] : self.range[1]]
+
+        if metrics.dim() == 2:
+            B, S = metrics.shape
+        elif metrics.dim() == 3:
+            B, S, _ = metrics.shape
+        else:
+            raise ValueError("metrics must be (B, S) or (B, S, K)")
+
+        for i in range(B):
+            rollout_id = current_paths[i]
+            df = self.tally.get(rollout_id, pd.DataFrame())
+
+            valid_idx = torch.nonzero(
+                self.action_mask[i].bool(), as_tuple=False
+            ).squeeze(-1)
+            if valid_idx.numel() == 0:
+                self.tally[rollout_id] = df
+                continue
+
+            idx_list = valid_idx.detach().cpu().tolist()
+
+            # Ensure index contains valid steps
+            if df.empty:
+                df = pd.DataFrame(index=idx_list)
+            else:
+                new_index = sorted(set(df.index.tolist()) | set(idx_list))
+                if list(df.index) != new_index:
+                    df = df.reindex(new_index)
+
+            # Slice metrics at valid steps
+            m_valid = metrics[i][valid_idx]
+
+            # -> pure python lists (1D list or list-of-lists)
+            values = m_valid.detach().cpu().tolist()
+
+            # optional tids -> tokens
+            if to_tids:
+
+                def _to_tokish(x):
+                    if isinstance(x, list):
+                        return self.tids_to_str([int(v) for v in x])
+                    else:
+                        return self.tids_to_str([int(x)])[0]
+
+                values = [_to_tokish(v) for v in values]
+
+            # Ensure column exists with object dtype, then assign via aligned Series
+            if metric_id not in df.columns:
+                df[metric_id] = pd.Series(index=df.index, dtype=object)
+
+            if isinstance(values, np.ndarray):
+                values = values.tolist()
+
+            if len(values) != len(idx_list):
+                raise ValueError(
+                    f"Length mismatch for '{metric_id}': values={len(values)} vs idx_list={len(idx_list)}"
+                )
+
+            df.loc[idx_list, metric_id] = pd.Series(
+                values, index=idx_list, dtype=object
+            )
+            self.tally[rollout_id] = df
+
+    # --------- Saving ---------
+
+    def save(self, path: str):
+        """
+        Write a manifest JSON and one CSV per rollout.
+
+        - Manifest includes metadata only (safe to JSON).
+        - Each rollout CSV is written with index label 'timestep'.
+        - Only a single 'context' column (list[str]).
+        """
+        if not self.tally or all(df.empty for df in self.tally.values()):
+            return
+
+        os.makedirs(path, exist_ok=True)
+        from datetime import datetime
+
+        now = datetime.now()
+
+        manifest = {
+            "created_at": f"{now:%Y-%m-%d %H:%M:%S}",
+            "max_context_length": self.max_context_length,
+            "num_rollouts": len(self.tally),
+            "rollouts": [],
+        }
+
+        for rid, df in self.tally.items():
+            rid_str = str(rid)
+            safe_name = self._sanitize_filename(rid_str)
+            csv_path = os.path.join(path, f"{safe_name}_tokenwise.csv")
+
+            # Put 'context' first, then the rest
+            cols = ["context"] + [c for c in df.columns if c != "context"]
+            try:
+                df[cols].to_csv(csv_path, index=True, index_label="timestep")
+            except Exception as e:
+                continue
+
+            manifest["rollouts"].append(
+                {
+                    "rollout_id": rid_str,
+                    "csv": csv_path,
+                    "num_rows": int(df.shape[0]),
+                    "columns": cols,
+                }
+            )
+
+        manifest_path = os.path.join(
+            path, f"tokenwise_manifest_{now:%Y-%m-%d___%H-%M-%S}.json"
+        )
+        with open(manifest_path, "w") as fp:
+            json.dump(manifest, fp, indent=2)

src_code_for_reproducibility/training/tokenize_chats.py

@@ -0,0 +1,128 @@
+"""
+File: mllm/training/tokenize_chats.py
+Summary: Tokenizes chat datasets and prepares tensors for training.
+"""
+
+import logging
+import sys
+
+import regex
+import torch
+from transformers import AutoTokenizer
+
+from mllm.training.training_data_utils import TrainingChatTurn, TrajectoryBatch
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def process_training_chat(
+    tokenizer: AutoTokenizer,
+    chat_history: list[TrainingChatTurn],
+    entropy_mask_regex: str | None = None,
+    exploration_prompts_to_remove: list[str] = [],
+    use_engine_out_token_ids: bool = False,
+) -> tuple[torch.IntTensor, torch.BoolTensor, torch.IntTensor, torch.BoolTensor]:
+    """Tokenize a single training chat and build aligned per-token masks.
+
+    Given an ordered list of `TrainingChatTurn`, this function tokenizes each
+    turn independently using the tokenizer's chat template, then concatenates
+    all resulting token sequences. It also constructs three parallel 1D masks
+    that align with the concatenated tokens:
+
+    - input_ids: token ids for the entire chat, turn by turn
+    - action_mask: True for tokens that belong to assistant turns (i.e., model
+      actions), False for tokens from other roles
+    - timesteps: per-token time step copied from the originating turn's
+      `time_step`
+    - state_ends_mask: True for the last token of any turn where
+      `is_state_end` is True, otherwise False
+
+    Important details:
+    - Each turn is passed as a single-message list to
+      `tokenizer.apply_chat_template` and flattened; the per-turn outputs are
+      then concatenated in the original order.
+    - Turn boundaries are not explicitly encoded beyond what the chat template
+      inserts; masks provide alignment for learning signals and state endings.
+    - No truncation or padding is performed here; downstream code should handle
+      batching/padding as needed.
+    - Note on dtypes: `input_ids` will be a LongTensor (int64). `action_mask`
+      and `state_ends_mask` are BoolTensors. `timesteps` is currently created
+      as a float tensor; adjust the implementation if integer dtype is
+      required downstream.
+
+    Args:
+        tokenizer: A Hugging Face tokenizer supporting `apply_chat_template`.
+        chat_history: Ordered list of `TrainingChatTurn` forming one dialogue.
+
+    Returns:
+        A tuple of four 1D tensors, all of equal length N (the total number of
+        tokens across all turns), in the following order:
+        - input_ids (LongTensor)
+        - action_mask (BoolTensor)
+        - timesteps (FloatTensor as implemented; see note above)
+        - state_ends_mask (BoolTensor)
+    """
+    state_ends_mask = []
+    input_ids = []
+    action_mask = []
+    timesteps = []
+    entropy_mask = []
+    engine_log_probs = []
+    for train_chat_turn in chat_history:
+        is_state_end = train_chat_turn.is_state_end
+        time_step = train_chat_turn.time_step
+        is_action = train_chat_turn.role == "assistant"
+
+        # Remove exploration prompts from training data
+        for exploration_prompt in exploration_prompts_to_remove:
+            if exploration_prompt in train_chat_turn.content:
+                train_chat_turn.content = train_chat_turn.content.replace(
+                    exploration_prompt, ""
+                )
+
+        chat_turn = {
+            "role": train_chat_turn.role,
+            "content": train_chat_turn.content,
+        }
+        if entropy_mask_regex is not None:
+            is_entropy_mask_true = (
+                regex.search(entropy_mask_regex, train_chat_turn.content) is not None
+            )
+        else:
+            is_entropy_mask_true = True
+        if is_action:
+            chat_turn_ids = train_chat_turn.out_token_ids
+            nb_chat_turns_ids = chat_turn_ids.numel()
+            action_mask.append(torch.ones(nb_chat_turns_ids, dtype=torch.bool))
+            engine_log_probs.append(train_chat_turn.log_probs)
+        else:
+            chat_turn_ids = train_chat_turn.chat_template_token_ids
+            nb_chat_turns_ids = chat_turn_ids.numel()
+            action_mask.append(torch.zeros(nb_chat_turns_ids, dtype=torch.bool))
+            engine_log_probs.append(torch.zeros(nb_chat_turns_ids, dtype=torch.float))
+        nb_chat_turns_ids = chat_turn_ids.numel()
+        state_ends_mask.append(torch.zeros(nb_chat_turns_ids, dtype=torch.bool))
+        if is_state_end:
+            state_ends_mask[-1][-1] = True  # last token is state end
+        input_ids.append(chat_turn_ids)
+        entropy_mask.append(torch.ones(nb_chat_turns_ids, dtype=torch.bool))
+        if not is_entropy_mask_true:
+            entropy_mask[-1] = entropy_mask[-1] * False
+        timesteps.append(torch.ones(nb_chat_turns_ids) * time_step)
+    input_ids = torch.cat(input_ids)
+    action_mask = torch.cat(action_mask)
+    entropy_mask = torch.cat(entropy_mask)
+    timesteps = torch.cat(timesteps)
+    timesteps = timesteps.to(torch.long)
+    state_ends_mask = torch.cat(state_ends_mask)
+    engine_log_probs = torch.cat(engine_log_probs)
+
+    return (
+        input_ids,
+        action_mask,
+        entropy_mask,
+        timesteps,
+        state_ends_mask,
+        engine_log_probs,
+    )

src_code_for_reproducibility/training/trainer_ad_align.py

@@ -0,0 +1,505 @@
+"""
+File: mllm/training/trainer_ad_align.py
+Summary: Trainer specialized for the advantage-alignment objective.
+"""
+
+import copy
+import logging
+import sys
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from mllm.markov_games.rollout_tree import (
+    ChatTurn,
+    RolloutTreeBranchNode,
+    RolloutTreeRootNode,
+)
+from mllm.training.credit_methods import (
+    get_advantage_alignment_credits,
+    get_discounted_state_visitation_credits,
+)
+from mllm.training.tally_metrics import Tally
+from mllm.training.tally_rollout import RolloutTally, RolloutTallyItem
+from mllm.training.tally_tokenwise import ContextualizedTokenwiseTally
+from mllm.training.tokenize_chats import process_training_chat
+from mllm.training.trainer_common import BaseTrainer
+from mllm.training.training_data_utils import (
+    AdvantagePacket,
+    TrainingBatch,
+    TrainingChatTurn,
+    TrajectoryBatch,
+    get_main_chat_list_and_rewards,
+    get_tokenwise_credits,
+)
+from mllm.utils.resource_context import resource_logger_context
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+RolloutId = int
+AgentId = str
+
+
+@dataclass
+class AdAlignTrainingData:
+    """Holds tensorized rollouts plus precomputed advantages for one agent."""
+
+    agent_id: str
+    main_data: TrajectoryBatch
+    # list-of-tensors: per rollout advantages with length jT
+    main_advantages: list[torch.FloatTensor] | None = None
+    # list-of-tensors: per rollout matrix (jT, A)
+    alternative_advantages: list[torch.FloatTensor] | None = None
+    advantage_alignment_credits: list[torch.FloatTensor] | None = None
+
+
+def get_alternative_chat_histories(
+    agent_id: str, root: RolloutTreeRootNode
+) -> list[list[TrainingChatTurn], list[torch.FloatTensor]]:
+    """
+    Traverse every unilateral branch under ``root`` and collect chat/reward histories.
+
+    Returns
+    -------
+    alternative_chats:
+        Flattened list of chat turns for each branch (ordered by branch depth).
+    alternative_rewards:
+        Matching list of reward tensors aligned with the chat history.
+    """
+    current_node = root.child
+    branches = current_node.branches
+    pre_branch_chat = []
+    pre_branch_rewards = []
+    alternative_rewards = []
+    alternative_chats = []
+    while current_node is not None:
+        assert isinstance(
+            current_node, RolloutTreeBranchNode
+        ), "Current node should be a branch node."
+        main_node = current_node.main_child
+        branches = current_node.branches
+        current_node = main_node.child
+
+        # Get the `A` alternative trajectories
+        alternative_nodes = branches[agent_id]
+        for alt_node in alternative_nodes:
+            post_branch_chat, post_branch_rewards = get_main_chat_list_and_rewards(
+                agent_id=agent_id, root=alt_node
+            )
+            branch_chat = pre_branch_chat + post_branch_chat
+            alternative_chats.append(branch_chat)
+            alternative_rewards.append(
+                torch.cat([torch.tensor(pre_branch_rewards), post_branch_rewards])
+            )
+
+        chat_turns: list[ChatTurn] = main_node.step_log.action_logs[agent_id].chat_turns
+        chat_turns: list[TrainingChatTurn] = [
+            TrainingChatTurn(time_step=main_node.time_step, **turn.model_dump())
+            for turn in chat_turns
+        ]
+
+        pre_branch_chat.extend(chat_turns)
+        pre_branch_rewards.append(
+            main_node.step_log.simulation_step_log.rewards[agent_id]
+        )
+
+    return alternative_chats, alternative_rewards
+
+
+class TrainerAdAlign(BaseTrainer):
+    """
+    Extends the reinforce trainer to support Advantage Alignment.
+    """
+
+    def __init__(
+        self,
+        ad_align_beta: float,
+        ad_align_gamma: float,
+        ad_align_exclude_k_equals_t: bool,
+        ad_align_use_sign: bool,
+        ad_align_clipping: float,
+        ad_align_force_coop_first_step: bool,
+        use_old_ad_align: bool,
+        use_time_regularization: bool,
+        rloo_branch: bool,
+        reuse_baseline: bool,
+        ad_align_beta_anneal_step: int = -1,
+        ad_align_beta_anneal_rate: float = 0.5,
+        min_ad_align_beta: float = 0.1,
+        mean_normalize_ad_align: bool = False,
+        whiten_adalign_advantages: bool = False,
+        whiten_adalign_advantages_time_step_wise: bool = False,
+        ad_align_discount_t: bool = False,
+        *args,
+        **kwargs,
+    ):
+        """
+        Initialize the advantage alignment trainer.
+        Args:
+            ad_align_beta: Beta parameter for the advantage alignment.
+            ad_align_gamma: Gamma parameter for the advantage alignment.
+            ad_align_exclude_k_equals_t: Whether to include k = t in the advantage alignment.
+            ad_align_use_sign: Whether to use sign in the advantage alignment.
+            ad_align_clipping: Clipping value for the advantage alignment.
+            ad_align_force_coop_first_step: Whether to force coop on the first step of the advantage alignment.
+        """
+        super().__init__(*args, **kwargs)
+        self.ad_align_beta = ad_align_beta
+        self.ad_align_gamma = ad_align_gamma
+        self.ad_align_exclude_k_equals_t = ad_align_exclude_k_equals_t
+        self.ad_align_use_sign = ad_align_use_sign
+        self.ad_align_clipping = ad_align_clipping
+        self.ad_align_force_coop_first_step = ad_align_force_coop_first_step
+        self.use_old_ad_align = use_old_ad_align
+        self.use_time_regularization = use_time_regularization
+        self.rloo_branch = rloo_branch
+        self.reuse_baseline = reuse_baseline
+        self.ad_align_beta_anneal_step = ad_align_beta_anneal_step
+        self.ad_align_beta_anneal_rate = ad_align_beta_anneal_rate
+        self.min_ad_align_beta = min_ad_align_beta
+        self.past_ad_align_step = -1
+        self.mean_normalize_ad_align = mean_normalize_ad_align
+        self.whiten_adalign_advantages = whiten_adalign_advantages
+        self.whiten_adalign_advantages_time_step_wise = (
+            whiten_adalign_advantages_time_step_wise
+        )
+        self.ad_align_discount_t = ad_align_discount_t
+        self.training_data: dict[AgentId, AdAlignTrainingData] = {}
+        self.debug_path_list: list[str] = []
+
+    def set_agent_trajectory_data(
+        self, agent_id: str, roots: list[RolloutTreeRootNode]
+    ):
+        """
+        Materialize main and alternative trajectory tensors used by the advantage-alignment trainer.
+        """
+
+        B = len(roots)  # Number of rollouts
+
+        # For main rollouts
+        batch_rollout_ids = []
+        batch_crn_ids = []
+        batch_input_ids = []
+        batch_action_mask = []
+        batch_entropy_mask = []
+        batch_timesteps = []
+        batch_state_ends_mask = []
+        batch_engine_log_probs = []
+        batch_rewards = []
+
+        # For alternative actions rollouts
+        batch_branching_time_steps = []
+        alternative_batch_input_ids = []
+        alternative_batch_action_mask = []
+        alternative_batch_entropy_mask = []
+        alternative_batch_timesteps = []
+        alternative_batch_state_ends_mask = []
+        alternative_batch_engine_log_probs = []
+        alternative_batch_rewards = []
+        jT_list = []
+
+        try:
+            A = len(roots[0].child.branches[agent_id])  # Number of alternative actions
+        except:
+            A = 0
+
+        for root in roots:
+            rollout_id = root.id
+            self.debug_path_list.append(
+                "mgid:" + str(rollout_id) + "_agent_id:" + agent_id
+            )
+            # Get main trajectory
+            batch_rollout_ids.append(rollout_id)
+            batch_crn_ids.append(root.crn_id)
+            main_chat, main_rewards = get_main_chat_list_and_rewards(
+                agent_id=agent_id, root=root
+            )
+            (
+                input_ids,
+                action_mask,
+                entropy_mask,
+                timesteps,
+                state_ends_mask,
+                engine_log_probs,
+            ) = process_training_chat(
+                tokenizer=self.tokenizer,
+                chat_history=main_chat,
+                entropy_mask_regex=self.entropy_mask_regex,
+                exploration_prompts_to_remove=self.exploration_prompts_to_remove,
+            )
+            batch_input_ids.append(input_ids)
+            batch_action_mask.append(action_mask)
+            batch_entropy_mask.append(entropy_mask)
+            batch_timesteps.append(timesteps)
+            batch_state_ends_mask.append(state_ends_mask)
+            batch_engine_log_probs.append(engine_log_probs)
+            batch_rewards.append(main_rewards)
+            jT = (
+                main_rewards.numel()
+            )  # Number of timesteps inferred from reward tensor length.
+            jT_list.append(jT)
+            if A > 0:
+                # We get the branching time steps for each of the `jT` time steps in the main trajectory.
+                branching_time_steps = [bt for item in range(jT) for bt in A * [item]]
+                batch_branching_time_steps.extend(branching_time_steps)
+
+                # Get all of the (jT*A) alternative trajectories in the tree
+                # (jT is the number of time steps in the main trajectory, A is the number of alternative actions)
+                alternative_chats, alternative_rewards = get_alternative_chat_histories(
+                    agent_id=agent_id, root=root
+                )
+                assert (
+                    len(alternative_chats) == A * jT
+                ), "Incorrect number of alternative trajectories."
+
+                for chat, rewards in zip(alternative_chats, alternative_rewards):
+                    (
+                        input_ids,
+                        action_mask,
+                        entropy_mask,
+                        timesteps,
+                        state_ends_mask,
+                        engine_log_probs,
+                    ) = process_training_chat(
+                        tokenizer=self.tokenizer,
+                        chat_history=chat,
+                        entropy_mask_regex=self.entropy_mask_regex,
+                        exploration_prompts_to_remove=self.exploration_prompts_to_remove,
+                    )
+                    alternative_batch_input_ids.append(input_ids)
+                    alternative_batch_action_mask.append(action_mask)
+                    alternative_batch_entropy_mask.append(entropy_mask)
+                    alternative_batch_timesteps.append(timesteps)
+                    alternative_batch_state_ends_mask.append(state_ends_mask)
+                    alternative_batch_engine_log_probs.append(engine_log_probs)
+                    alternative_batch_rewards.append(rewards)
+
+        jT_list = torch.Tensor(jT_list)
+
+        # Assert that number of alternative actions is constant
+        # assert len(set(nb_alternative_actions)) == 1, "Number of alternative actions must be constant"
+        # A = nb_alternative_actions[0]
+
+        trajectory_batch = TrajectoryBatch(
+            rollout_ids=torch.tensor(batch_rollout_ids, dtype=torch.int32),  # (B,)
+            crn_ids=torch.tensor(batch_crn_ids, dtype=torch.int32),
+            agent_ids=[agent_id] * len(batch_rollout_ids),
+            batch_input_ids=batch_input_ids,
+            batch_action_mask=batch_action_mask,
+            batch_entropy_mask=batch_entropy_mask,
+            batch_timesteps=batch_timesteps,
+            batch_state_ends_mask=batch_state_ends_mask,
+            batch_engine_log_probs=batch_engine_log_probs,
+            batch_rewards=batch_rewards,
+        )
+        # Get Advantages & Train Critic
+        with resource_logger_context(
+            logger, "Get advantages with critic gradient accumulation"
+        ):
+            self.batch_advantages: torch.FloatTensor = (
+                self.get_advantages_with_critic_gradient_accumulation(trajectory_batch)
+            )  # (B, jT)
+
+        if A > 0:
+            # Here, `A` is the number of alternative actions / trajectories taken at each time step.
+            # For each of the `B` rollout perspectives, at each of its jT (`j` is for jagged, since each main rollout may be of a different length) steps, we take A alternate trajectories (from different actions).
+            # Therefore, we have ∑jT * A trajectories to process. If each of the main trajectories have T steps, we will have `B*T*A` to process.
+            with resource_logger_context(logger, "Create alternative trajectory batch"):
+                sum_jT = int(torch.sum(jT_list).item())
+                jT_list = (
+                    jT_list.int().tolist()
+                )  # (jT,) # (we only want the advantages where we branched out)
+                alternative_trajectory_batch = TrajectoryBatch(
+                    rollout_ids=torch.zeros(A * sum_jT, dtype=torch.int32),
+                    crn_ids=torch.zeros(A * sum_jT, dtype=torch.int32),
+                    agent_ids=[agent_id] * (A * sum_jT),
+                    batch_input_ids=alternative_batch_input_ids,
+                    batch_action_mask=alternative_batch_action_mask,
+                    batch_entropy_mask=alternative_batch_entropy_mask,
+                    batch_timesteps=alternative_batch_timesteps,
+                    batch_state_ends_mask=alternative_batch_state_ends_mask,
+                    batch_engine_log_probs=alternative_batch_engine_log_probs,
+                    batch_rewards=alternative_batch_rewards,
+                )
+
+            # Get alternative advantages
+            # BAAs stands for batch alternative advantages
+            # (torch nested tensors have very little api support, so we have to do some odd manual work here)
+            with resource_logger_context(
+                logger, "Compute alternative advantage estimates"
+            ):
+                BAAs_list = self.get_advantages_with_critic_gradient_accumulation(
+                    alternative_trajectory_batch
+                )  # list length (∑jT * A), each (jT',)
+                # Pad alternative advantages to (∑jT*A, P)
+
+                BAAs_padded = pad_sequence(
+                    BAAs_list, batch_first=True, padding_value=0.0
+                )
+                branch_idx = torch.tensor(
+                    batch_branching_time_steps,
+                    device=BAAs_padded.device,
+                    dtype=torch.long,
+                )
+                gathered = BAAs_padded.gather(
+                    dim=1, index=branch_idx.unsqueeze(1)
+                ).squeeze(1)
+                # Reshape and split per rollout, then transpose to (jT_i, A)
+                gathered = gathered.view(A, sum_jT)  # (A, ∑jT)
+                blocks = list(
+                    torch.split(gathered, jT_list, dim=1)
+                )  # len B, shapes (A, jT_i)
+                BAAs = [
+                    blk.transpose(0, 1).contiguous() for blk in blocks
+                ]  # list of (jT_i, A)
+        if self.ad_align_beta_anneal_step > 0:
+            max_rollout_id = torch.max(trajectory_batch.rollout_ids) + 1
+            if (
+                max_rollout_id % self.ad_align_beta_anneal_step == 0
+                and self.past_ad_align_step != max_rollout_id
+            ):
+                self.ad_align_beta = max(
+                    self.ad_align_beta * self.ad_align_beta_anneal_rate,
+                    self.min_ad_align_beta,
+                )
+                logger.info(f"Annealing ad_align_beta to {self.ad_align_beta}")
+                self.past_ad_align_step = max_rollout_id
+        self.training_data[agent_id] = AdAlignTrainingData(
+            agent_id=agent_id,
+            main_data=trajectory_batch,
+            main_advantages=self.batch_advantages,
+            alternative_advantages=BAAs if A > 0 else None,
+        )
+
+    def share_advantage_data(self) -> list[AdvantagePacket]:
+        """
+        Share the advantage alignment data with other agents.
+        Returns:
+            AdvantagePacket: The advantage packet containing the agent's advantages.
+        """
+        logger.info(f"Sharing advantage alignment data.")
+        advantage_packets = []
+        for _, agent_data in self.training_data.items():
+            advantage_packets.append(
+                AdvantagePacket(
+                    agent_id=agent_data.agent_id,
+                    rollout_ids=agent_data.main_data.rollout_ids,
+                    main_advantages=agent_data.main_advantages,
+                )
+            )
+        return advantage_packets
+
+    def receive_advantage_data(self, advantage_packets: list[AdvantagePacket]):
+        """
+        Receive advantage packets from other players.
+        These contain the advantages of the other players' rollouts estimated by them.
+        """
+        logger.info(f"Receiving advantage packets.")
+
+        assert (
+            len(advantage_packets) > 0
+        ), "At least one advantage packet must be provided."
+
+        for agent_id, agent_data in self.training_data.items():
+            coagent_advantage_packets = [
+                packet for packet in advantage_packets if packet.agent_id != agent_id
+            ]
+            agent_rollout_ids = agent_data.main_data.rollout_ids
+            agent_advantages = agent_data.main_advantages
+            co_agent_advantages = []
+            for rollout_id in agent_rollout_ids:
+                for co_agent_packet in coagent_advantage_packets:
+                    if rollout_id in co_agent_packet.rollout_ids:
+                        index = torch.where(rollout_id == co_agent_packet.rollout_ids)[
+                            0
+                        ].item()
+                        co_agent_advantages.append(
+                            co_agent_packet.main_advantages[index]
+                        )
+                        # assumes that its two player game, with one co-agent
+                        break
+            assert len(co_agent_advantages) == len(agent_advantages)
+            B = len(agent_advantages)
+            assert all(
+                a.shape[0] == b.shape[0]
+                for a, b in zip(co_agent_advantages, agent_advantages)
+            ), "Number of advantages must match for advantage alignment."
+
+            # Get padded tensors (advantage alignment is invariant to padding)
+            lengths = torch.tensor(
+                [len(t) for t in agent_advantages],
+                device=self.device,
+                dtype=torch.long,
+            )
+            padded_main_advantages = pad_sequence(
+                agent_advantages, batch_first=True, padding_value=0.0
+            )
+            if agent_data.alternative_advantages:
+                padded_alternative_advantages = pad_sequence(
+                    agent_data.alternative_advantages,
+                    batch_first=True,
+                    padding_value=0.0,
+                )  # (B, P, A)
+            else:
+                padded_alternative_advantages = None
+            padded_co_agent_advantages = pad_sequence(
+                co_agent_advantages, batch_first=True, padding_value=0.0
+            )
+
+            # Create training batch data
+            credits, sub_tensors = get_advantage_alignment_credits(
+                a1=padded_main_advantages,
+                a1_alternative=padded_alternative_advantages,
+                a2=padded_co_agent_advantages,
+                beta=self.ad_align_beta,
+                gamma=self.ad_align_gamma,
+                exclude_k_equals_t=self.ad_align_exclude_k_equals_t,
+                use_sign=self.ad_align_use_sign,
+                clipping=self.ad_align_clipping,
+                force_coop_first_step=self.ad_align_force_coop_first_step,
+                use_old_ad_align=self.use_old_ad_align,
+                use_time_regularization=self.use_time_regularization,
+                rloo_branch=self.rloo_branch,
+                reuse_baseline=self.reuse_baseline,
+                mean_normalize_ad_align=self.mean_normalize_ad_align,
+                whiten_adalign_advantages=self.whiten_adalign_advantages,
+                whiten_adalign_advantages_time_step_wise=self.whiten_adalign_advantages_time_step_wise,
+                discount_t=self.ad_align_discount_t,
+            )
+            for key, value in sub_tensors.items():
+                self.rollout_tally.add_metric(
+                    path=[key],
+                    rollout_tally_item=RolloutTallyItem(
+                        crn_ids=agent_data.main_data.crn_ids,
+                        rollout_ids=agent_data.main_data.rollout_ids,
+                        agent_ids=agent_data.main_data.agent_ids,
+                        metric_matrix=value,
+                    ),
+                )
+
+            if not self.skip_discounted_state_visitation:
+                credits = get_discounted_state_visitation_credits(
+                    credits,
+                    self.discount_factor,
+                )
+                self.rollout_tally.add_metric(
+                    path=["discounted_state_visitation_credits"],
+                    rollout_tally_item=RolloutTallyItem(
+                        crn_ids=agent_data.main_data.crn_ids,
+                        rollout_ids=agent_data.main_data.rollout_ids,
+                        agent_ids=agent_data.main_data.agent_ids,
+                        metric_matrix=sub_tensors[
+                            "discounted_state_visitation_credits"
+                        ],
+                    ),
+                )
+
+            # Slice back to jagged
+            advantage_alignment_credits = [credits[i, : lengths[i]] for i in range(B)]
+            # Replace stored training data for this agent by the concrete trajectory batch
+            # and attach the computed credits for policy gradient.
+            self.training_data[agent_id] = agent_data.main_data
+            self.training_data[agent_id].batch_credits = advantage_alignment_credits