Add files using upload-large-folder tool

.hydra/config.yaml

@@ -0,0 +1,183 @@
+experiment:
+  wandb_enabled: true
+  nb_epochs: 3000
+  nb_matches_per_iteration: 64
+  reinit_matches_each_it: true
+  checkpoint_every_n_iterations: 50
+  start_epoch: 0
+  resume_experiment: true
+  base_seed: 0
+  seed_group_size: 8
+  train: true
+  stat_methods_for_live_wandb: mllm.markov_games.negotiation.negotiation_statistics
+  name: naive_vs_fixed_ad_align_seed4321
+  agent_buffer: false
+  keep_agent_buffer_count: ${lora_count}
+  agent_buffer_recent_k: -1
+  description: Trust-and-Split Rock Paper Scissors negotiation game
+logging:
+  wandb:
+    enabled: false
+    project: llm-negotiation
+    entity: null
+    mode: online
+    name: null
+    group: null
+    tags: []
+    notes: null
+temperature: 1.0
+markov_games:
+  runner_method_name: LinearRunner
+  runner_kwargs: {}
+  group_by_round: true
+  simulation_class_name: TrustAndSplitRPSSimulation
+  simulation_init_args:
+    nb_of_rounds: 10
+    quota_messages_per_agent_per_round: 1
+    alternating_hands: false
+  agents:
+    0:
+      agent_id: ${agent_0_id}
+      agent_name: Alice
+      agent_class_name: TrustAndSplitRPSAgent
+      policy_id: base_llm/agent_adapter
+      init_kwargs:
+        goal: Maximize your total points over the whole game.
+        num_message_chars: 500
+        message_start_end_format: true
+        proposal_start_end_format: true
+    1:
+      agent_id: ${agent_1_id}
+      agent_name: Bob
+      agent_class_name: TrustAndSplitRPSAgent
+      policy_id: base_llm/fixed_ad_align_adapter
+      init_kwargs:
+        goal: Maximize your total points over the whole game.
+        num_message_chars: 500
+        message_start_end_format: true
+        proposal_start_end_format: true
+models:
+  base_llm:
+    class: LeanLocalLLM
+    init_args:
+      llm_id: base_llm
+      model_name: Qwen/Qwen2.5-7B-Instruct
+      inference_backend: vllm
+      hf_kwargs:
+        device_map: auto
+        torch_dtype: bfloat16
+        max_memory:
+          0: 20GiB
+        attn_implementation: flash_attention_2
+      inference_backend_init_kwargs:
+        enable_lora: true
+        seed: ${experiment.base_seed}
+        enable_prefix_caching: true
+        max_model_len: 10000.0
+        gpu_memory_utilization: 0.5
+        dtype: bfloat16
+        trust_remote_code: true
+        max_lora_rank: 32
+        enforce_eager: false
+        max_loras: ${lora_count}
+        max_cpu_loras: ${lora_count}
+        enable_sleep_mode: true
+      inference_backend_sampling_params:
+        temperature: ${temperature}
+        top_p: 1.0
+        max_tokens: 400
+        top_k: -1
+        logprobs: 0
+      adapter_configs:
+        agent_adapter:
+          task_type: CAUSAL_LM
+          r: 32
+          lora_alpha: 64
+          lora_dropout: 0.0
+          target_modules: all-linear
+        critic_adapter:
+          task_type: CAUSAL_LM
+          r: 32
+          lora_alpha: 64
+          lora_dropout: 0.0
+          target_modules: all-linear
+        fixed_ad_align_adapter:
+          task_type: CAUSAL_LM
+          r: 32
+          lora_alpha: 64
+          lora_dropout: 0.0
+          target_modules: all-linear
+      enable_thinking: null
+      regex_max_attempts: 1
+      initial_adapter_paths:
+        fixed_ad_align_adapter: ${fixed_ad_align_adapter_path}
+critics:
+  agent_critic:
+    module_pointer:
+    - base_llm
+    - critic_adapter
+optimizers:
+  agent_optimizer:
+    module_pointer:
+    - base_llm
+    - agent_adapter
+    optimizer_class_name: torch.optim.Adam
+    init_args:
+      lr: 3.0e-06
+      weight_decay: 0.0
+  critic_optimizer:
+    module_pointer: agent_critic
+    optimizer_class_name: torch.optim.Adam
+    init_args:
+      lr: 3.0e-06
+      weight_decay: 0.0
+trainers:
+  agent_trainer:
+    class: TrainerNaive
+    module_pointers:
+      policy:
+      - base_llm
+      - agent_adapter
+      policy_optimizer: agent_optimizer
+      critic: agent_critic
+      critic_optimizer: critic_optimizer
+    kwargs:
+      entropy_coeff: 0.0
+      entropy_topk: null
+      entropy_mask_regex: null
+      kl_coeff: 0.001
+      gradient_clipping: 1.0
+      restrict_tokens: null
+      mini_batch_size: 1
+      use_gradient_checkpointing: true
+      temperature: ${temperature}
+      device: cuda:0
+      use_gae: false
+      whiten_advantages: false
+      whiten_advantages_time_step_wise: false
+      skip_discounted_state_visitation: true
+      use_gae_lambda_annealing: false
+      gae_lambda_annealing_method: None
+      gae_lambda_annealing_method_params: None
+      gae_lambda_annealing_limit: 0.95
+      discount_factor: 0.96
+      use_rloo: true
+      enable_tokenwise_logging: false
+      pg_loss_normalization: nb_tokens
+      truncated_importance_sampling_ratio_cap: 2.0
+      reward_normalizing_constant: 100.0
+train_on_which_data:
+  agent_trainer:
+  - Alice
+lora_count: 30
+common_agent_kwargs:
+  goal: Maximize your total points over the whole game.
+  num_message_chars: 500
+  message_start_end_format: true
+  proposal_start_end_format: true
+agent_0_id: Alice
+agent_1_id: Bob
+agent_ids:
+- Alice
+- Bob
+fixed_ad_align_adapter_path: /home/muqeeth/scratch/llm_negotiation/2025_11/tas_rps_startend_ad_align_nocurrtimestep_seed4321_beta2/seed_4321/Qwen/Qwen2.5-7B-Instruct/adapters/agent_adapter

.hydra/hydra.yaml

@@ -0,0 +1,154 @@
+hydra:
+  run:
+    dir: ${oc.env:SCRATCH}/llm_negotiation/${now:%Y_%m}/${experiment.name}
+  sweep:
+    dir: multirun/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.mode=RUN
+    task: []
+  job:
+    name: run
+    chdir: false
+    override_dirname: ''
+    id: ???
+    num: ???
+    config_name: naive_vs_fixed_ad_align_seed4321.yaml
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.1'
+    cwd: /scratch/muqeeth/llm_negotiation
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /scratch/muqeeth/llm_negotiation/configs
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /scratch/muqeeth/llm_negotiation/2025_11/naive_vs_fixed_ad_align_seed4321
+    choices:
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

seed_0/Qwen/Qwen2.5-7B-Instruct/adapters/README.md

@@ -0,0 +1,207 @@
+---
+base_model: Qwen/Qwen2.5-7B-Instruct
+library_name: peft
+pipeline_tag: text-generation
+tags:
+- base_model:adapter:Qwen/Qwen2.5-7B-Instruct
+- lora
+- transformers
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.17.1

seed_0/Qwen/Qwen2.5-7B-Instruct/adapters/agent_adapter/adapter_config.json

@@ -0,0 +1,42 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "Qwen/Qwen2.5-7B-Instruct",
+  "bias": "none",
+  "corda_config": null,
+  "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",
+  "qalora_group_size": 16,
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "v_proj",
+    "k_proj",
+    "q_proj",
+    "down_proj"
+  ],
+  "target_parameters": null,
+  "task_type": "CAUSAL_LM",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false
+}

seed_0/Qwen/Qwen2.5-7B-Instruct/adapters/critic_adapter/adapter_config.json

@@ -0,0 +1,42 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "Qwen/Qwen2.5-7B-Instruct",
+  "bias": "none",
+  "corda_config": null,
+  "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",
+  "qalora_group_size": 16,
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "v_proj",
+    "k_proj",
+    "q_proj",
+    "down_proj"
+  ],
+  "target_parameters": null,
+  "task_type": "CAUSAL_LM",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false
+}

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

@@ -0,0 +1,42 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "Qwen/Qwen2.5-7B-Instruct",
+  "bias": "none",
+  "corda_config": null,
+  "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",
+  "qalora_group_size": 16,
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "v_proj",
+    "k_proj",
+    "q_proj",
+    "down_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,84 @@
+import torch
+
+from mllm.chat_utils.chat_turn import ChatTurn
+from mllm.chat_utils.template_specific import (
+    custom_gemma3_template,
+    custom_llama3_template,
+    custom_qwen2_template,
+    custom_qwen3_template,
+    gemma3_assistant_postfix,
+    qwen2_assistant_postfix,
+    qwen3_assistant_postfix,
+)
+
+
+def get_custom_chat_template(tokenizer) -> str:
+    """
+    Get the chat template for the tokenizer.
+    """
+    if "qwen2" in tokenizer.name_or_path.lower():
+        return custom_qwen2_template
+    elif "llama" in tokenizer.name_or_path.lower():
+        return custom_llama3_template
+    elif "qwen3" in tokenizer.name_or_path.lower():
+        return custom_qwen3_template
+    elif "gemma" in tokenizer.name_or_path.lower():
+        return custom_gemma3_template
+    else:
+        raise ValueError(f"Tokenizer {tokenizer.name_or_path} not supported")
+
+
+def get_custom_assistant_postfix(tokenizer) -> torch.Tensor:
+    """
+    Get the custom assistant postfix for the tokenizer.
+    """
+    if "qwen2" in tokenizer.name_or_path.lower():
+        return qwen2_assistant_postfix
+    elif "qwen3" in tokenizer.name_or_path.lower():
+        return qwen3_assistant_postfix
+    elif "gemma" in tokenizer.name_or_path.lower():
+        return gemma3_assistant_postfix
+    return torch.tensor([], dtype=torch.long)
+
+
+def tokenize_chats(chats: list[ChatTurn], tokenizer, enable_thinking) -> None:
+    """
+    Set the chat_template_token_ids for each chat turn.
+    # TODO: use engine tokens if available
+    """
+    custom_template = get_custom_chat_template(tokenizer)
+    custom_assistant_postfix: torch.Tensor = get_custom_assistant_postfix(tokenizer)
+    for i, chat in enumerate(chats):
+        if chat.chat_template_token_ids is None:
+            if chat.role == "user":
+                next_chat = chats[i + 1] if i + 1 < len(chats) else None
+                add_generation_prompt = True
+                if next_chat and next_chat.role == "user":
+                    add_generation_prompt = False
+                encoded_chat = tokenizer.apply_chat_template(
+                    [chat],
+                    return_tensors="pt",
+                    chat_template=custom_template,
+                    add_generation_prompt=add_generation_prompt,
+                    add_system_prompt=True if i == 0 else False,
+                    enable_thinking=enable_thinking,
+                ).flatten()
+                previous_chat = chats[i - 1] if i > 0 else None
+                if previous_chat and previous_chat.role == "assistant":
+                    encoded_chat = torch.cat([custom_assistant_postfix, encoded_chat])
+            elif chat.role == "assistant":
+                encoded_chat = chat.out_token_ids
+            chat.chat_template_token_ids = encoded_chat
+
+
+def chat_turns_to_token_ids(
+    chats: list[ChatTurn], tokenizer, enable_thinking
+) -> list[int]:
+    """
+    Tokenize the chat turns and set the chat_template_token_ids for each chat turn.
+    """
+    tokenize_chats(chats=chats, tokenizer=tokenizer, enable_thinking=enable_thinking)
+    token_ids = []
+    for chat in chats:
+        token_ids.append(chat.chat_template_token_ids)
+    return torch.cat(token_ids)

src_code_for_reproducibility/chat_utils/chat_turn.py

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

src_code_for_reproducibility/chat_utils/template_specific.py

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

src_code_for_reproducibility/docs/Makefile

@@ -0,0 +1,19 @@
+# Minimal makefile for Sphinx documentation
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS)

src_code_for_reproducibility/docs/generate_docs.py

@@ -0,0 +1,249 @@
+#!/usr/bin/env python3
+"""
+Script to automatically generate Sphinx documentation for all modules and build the HTML website.
+"""
+import importlib.util
+import os
+import subprocess
+import sys
+
+
+def check_and_install_dependencies():
+    """Check for required dependencies and install them if missing."""
+    required_packages = [
+        "sphinx",
+        "sphinx-rtd-theme",
+        "sphinxcontrib-napoleon",
+        "sphinxcontrib-mermaid",
+        "sphinx-autodoc-typehints",
+    ]
+
+    missing_packages = []
+
+    for package in required_packages:
+        # Convert package name to module name (replace - with _)
+        module_name = package.replace("-", "_")
+
+        # Check if the package is installed
+        if importlib.util.find_spec(module_name) is None:
+            missing_packages.append(package)
+
+    # Install missing packages
+    if missing_packages:
+        print(f"Installing missing dependencies: {', '.join(missing_packages)}")
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install"] + missing_packages
+        )
+        print("Dependencies installed successfully")
+    else:
+        print("All required dependencies are already installed")
+
+
+def create_makefile(docs_dir):
+    """Create a Makefile for Sphinx documentation if it doesn't exist."""
+    makefile_path = os.path.join(docs_dir, "Makefile")
+
+    if os.path.exists(makefile_path):
+        print(f"Makefile already exists at {makefile_path}")
+        return
+
+    print(f"Creating Makefile at {makefile_path}")
+
+    makefile_content = """# Minimal makefile for Sphinx documentation
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(SPHINXFLAGS)
+"""
+
+    with open(makefile_path, "w") as f:
+        f.write(makefile_content)
+
+    print("Makefile created successfully")
+
+
+def create_make_bat(docs_dir):
+    """Create a make.bat file for Windows if it doesn't exist."""
+    make_bat_path = os.path.join(docs_dir, "make.bat")
+
+    if os.path.exists(make_bat_path):
+        print(f"make.bat already exists at {make_bat_path}")
+        return
+
+    print(f"Creating make.bat at {make_bat_path}")
+
+    make_bat_content = """@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd
+"""
+
+    with open(make_bat_path, "w") as f:
+        f.write(make_bat_content)
+
+    print("make.bat created successfully")
+
+
+def main():
+    # Check and install required dependencies
+    print("=== Checking dependencies ===")
+    check_and_install_dependencies()
+
+    # Get the directory of this script
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+
+    # Path to the project root
+    project_root = os.path.dirname(script_dir)
+
+    # Path to the source directory
+    source_dir = os.path.join(project_root, "src")
+
+    # Path to the docs source directory
+    docs_source_dir = os.path.join(script_dir, "source")
+
+    # Print paths for debugging
+    print(f"Script directory: {script_dir}")
+    print(f"Project root: {project_root}")
+    print(f"Source directory: {source_dir}")
+    print(f"Docs source directory: {docs_source_dir}")
+
+    # Make sure the source directory exists
+    if not os.path.exists(source_dir):
+        print(f"Error: Source directory {source_dir} does not exist!")
+        sys.exit(1)
+
+    # Make sure the docs source directory exists
+    if not os.path.exists(docs_source_dir):
+        print(f"Creating docs source directory: {docs_source_dir}")
+        os.makedirs(docs_source_dir)
+
+    # Step 1: Run sphinx-apidoc to generate .rst files for all modules
+    print("\n=== Generating API documentation ===")
+    cmd = [
+        "sphinx-apidoc",
+        "-f",  # Force overwriting of existing files
+        "-e",  # Put module documentation before submodule documentation
+        "-M",  # Put module documentation before subpackage documentation
+        "-o",
+        docs_source_dir,  # Output directory
+        source_dir,  # Source code directory
+    ]
+
+    print(f"Running command: {' '.join(cmd)}")
+    result = subprocess.run(cmd, capture_output=True, text=True)
+
+    # Print the output of the command
+    print("STDOUT:")
+    print(result.stdout)
+
+    print("STDERR:")
+    print(result.stderr)
+
+    if result.returncode != 0:
+        print(f"Error: sphinx-apidoc failed with return code {result.returncode}")
+        sys.exit(1)
+
+    # List the files in the docs source directory
+    print("\nFiles in docs/source directory:")
+    for file in sorted(os.listdir(docs_source_dir)):
+        print(f"  {file}")
+
+    print("\nDocumentation source files generated successfully!")
+
+    # Step 2: Create Makefile and make.bat if they don't exist
+    create_makefile(script_dir)
+    create_make_bat(script_dir)
+
+    # Step 3: Build the HTML documentation
+    print("\n=== Building HTML documentation ===")
+
+    # Determine the build command based on the platform
+    if os.name == "nt":  # Windows
+        build_cmd = ["make.bat", "html"]
+    else:  # Unix/Linux/Mac
+        build_cmd = ["make", "html"]
+
+    # Change to the docs directory to run the build command
+    os.chdir(script_dir)
+
+    print(f"Running command: {' '.join(build_cmd)}")
+    build_result = subprocess.run(build_cmd, capture_output=True, text=True)
+
+    # Print the output of the build command
+    print("STDOUT:")
+    print(build_result.stdout)
+
+    print("STDERR:")
+    print(build_result.stderr)
+
+    if build_result.returncode != 0:
+        print(f"Error: HTML build failed with return code {build_result.returncode}")
+        sys.exit(1)
+
+    # Get the path to the built HTML documentation
+    html_dir = os.path.join(script_dir, "build", "html")
+    index_path = os.path.join(html_dir, "index.html")
+
+    if os.path.exists(index_path):
+        print(f"\nHTML documentation built successfully!")
+        print(f"You can view it by opening: {index_path}")
+
+        # Try to open the documentation in a browser
+        try:
+            import webbrowser
+
+            print("\nAttempting to open documentation in your default browser...")
+            webbrowser.open(f"file://{index_path}")
+        except Exception as e:
+            print(f"Could not open browser automatically: {e}")
+    else:
+        print(f"\nWarning: HTML index file not found at {index_path}")
+
+
+if __name__ == "__main__":
+    main()

src_code_for_reproducibility/docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

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

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

src_code_for_reproducibility/docs/source/src.experiments.rst

@@ -0,0 +1,17 @@
+src.experiments package
+=======================
+
+.. automodule:: src.experiments
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Submodules
+----------
+
+.. toctree::
+   :maxdepth: 4
+
+   src.experiments.arithmetic_test
+   src.experiments.generate_and_train
+   src.experiments.last_completion

src_code_for_reproducibility/docs/source/src.generation.rst

@@ -0,0 +1,15 @@
+src.generation package
+======================
+
+.. automodule:: src.generation
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Submodules
+----------
+
+.. toctree::
+   :maxdepth: 4
+
+   src.generation.run_games

src_code_for_reproducibility/docs/source/src.models.rst

@@ -0,0 +1,20 @@
+src.models package
+==================
+
+.. automodule:: src.models
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Submodules
+----------
+
+.. toctree::
+   :maxdepth: 4
+
+   src.models.dummy_local_llm
+   src.models.local_llm
+   src.models.new_local_llm
+   src.models.server_llm
+   src.models.updatable_worker
+   src.models.vllm_worker_wrap

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

@@ -0,0 +1,7 @@
+src.training.ppo\_train\_value\_head module
+===========================================
+
+.. automodule:: src.training.ppo_train_value_head
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/markov_games/agent.py

@@ -0,0 +1,76 @@
+"""
+In simple RL paradise, where the action dimensions are constant and well defined,
+Agent classes are not necessary. But in MARL, with LLM's, there isn't always
+a direct path from policy to action. For instance, from the observation of the environment,
+a prompt must be created. Then, the outputs of the policy might be incorrect, so a second
+request to the LLM must be sent before the action is well defined. This is why this Agent class exists.
+It acts as a mini environment, bridging the gap between the core simulation and
+the LLM policies.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Tuple
+
+from numpy.random import default_rng
+
+from mllm.markov_games.rollout_tree import AgentActLog
+
+
+class Agent(ABC):
+    @abstractmethod
+    def __init__(
+        self,
+        seed: int,
+        agent_id: str,
+        agent_name: str,
+        agent_policy: Callable[[list[dict]], str],
+        *args,
+        **kwargs,
+    ):
+        """
+        Initialize the agent state.
+        """
+        self.seed = seed
+        self.agent_id = agent_id
+        self.agent_name = agent_name
+        self.policy = policy
+        self.rng = default_rng(self.seed)
+        raise NotImplementedError
+
+    async def act(self, observation) -> Tuple[Any, AgentActLog]:
+        """
+        Query (possibly multiple times) a policy (or possibly a pool of policies) to
+        obtain the action of the agent.
+
+        Example:
+        action = None
+        prompt = self.observation_to_prompt(observation)
+        while not self.valid(action):
+            output = await self.policy.generate(prompt)
+            action = self.policy_output_to_action(output)
+        return action
+
+        Returns:
+            action
+            step_info
+        """
+        raise NotImplementedError
+
+    def get_safe_copy(self):
+        """
+        Return copy of the agent object that is decorrelated from the original object.
+        """
+        raise NotImplementedError
+
+    def reset(self):
+        raise NotImplementedError
+
+    def render(self):
+        raise NotImplementedError
+
+    def close(self):
+        raise NotImplementedError
+
+    def get_agent_info(self):
+        raise NotImplementedError

src_code_for_reproducibility/markov_games/alternative_actions_runner.py

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

src_code_for_reproducibility/markov_games/group_timesteps.py

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

src_code_for_reproducibility/markov_games/linear_runner.py

@@ -0,0 +1,30 @@
+import asyncio
+import json
+import os.path
+
+from mllm.markov_games.markov_game import MarkovGame
+from mllm.markov_games.rollout_tree import RolloutTreeNode, RolloutTreeRootNode
+
+
+async def LinearRunner(
+    markov_game: MarkovGame, output_folder: str
+) -> RolloutTreeRootNode:
+    """
+    This method generates a trajectory without branching.
+    """
+    time_step = 0
+    terminated = False
+    root = RolloutTreeRootNode(
+        id=markov_game.get_id(),
+        crn_id=markov_game.get_crn_id(),
+        agent_ids=markov_game.get_agent_ids(),
+    )
+    previous_node = root
+    while not terminated:
+        terminated, step_log = await markov_game.step()
+        current_node = RolloutTreeNode(step_log=step_log, time_step=time_step)
+        previous_node.child = current_node
+        previous_node = current_node
+        time_step += 1
+
+    return root

src_code_for_reproducibility/markov_games/markov_game.py

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

src_code_for_reproducibility/markov_games/mg_utils.py

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

src_code_for_reproducibility/markov_games/rollout_tree.py

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

src_code_for_reproducibility/markov_games/run_markov_games.py

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

src_code_for_reproducibility/markov_games/simulation.py

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

src_code_for_reproducibility/markov_games/statistics_runner.py

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

src_code_for_reproducibility/markov_games/vine_ppo.py

@@ -0,0 +1,10 @@
+from anytree import Node, RenderTree
+from anytree.exporter import DotExporter
+import os.path
+import asyncio
+from mllm.markov_games.markov_game import MarkovGame
+
+async def VinePPORunner(
+    markov_game: MarkovGame,
+    **kwargs):
+    pass

src_code_for_reproducibility/models/adapter_training_wrapper.py

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

src_code_for_reproducibility/models/human_policy.py

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

src_code_for_reproducibility/models/inference_backend.py

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

src_code_for_reproducibility/models/inference_backend_dummy.py

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

src_code_for_reproducibility/models/inference_backend_sglang.py

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

src_code_for_reproducibility/models/inference_backend_sglang_local_server.py

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

src_code_for_reproducibility/models/inference_backend_vllm.py

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

src_code_for_reproducibility/models/inference_backend_vllm_local_server.py

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

src_code_for_reproducibility/models/large_language_model_api.py

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

src_code_for_reproducibility/models/large_language_model_local.py

@@ -0,0 +1,384 @@
+"""
+TODO: Figure out how to tweak SGlang not to go OOM when batch size is 32. See https://github.com/sgl-project/sglang/issues/6309.
+"""
+
+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 sglang.utils import (
+#     launch_server_cmd,
+#     print_highlight,
+#     terminate_process,
+#     wait_for_server,
+# )
+from torch.optim import SGD, Adam, AdamW, RMSprop
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import AutoModelForCausalLMWithValueHead
+
+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_sglang import SGLangOfflineBackend
+from mllm.models.inference_backend_vllm import VLLMAsyncBackend
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+AdapterID = str
+PolicyID = str
+
+
+class LeanLocalLLM:
+    """
+    TOWRITE
+    """
+
+    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", "sglang", "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 == "sglang":
+            self.inference_backend = SGLangOfflineBackend(
+                model_name=self.model_name,
+                save_path=self.save_path,
+                adapter_paths=self.adapter_paths,
+                tokenizer=self.tokenizer,
+                kwargs=inference_backend_init_kwargs,
+            )
+        elif 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]:
+        """
+        TOWRITE
+        """
+        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.
+        # TODO: create the adapter wrappers here
+        See adapter_wrapper.py
+        """
+        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,
+            )
+            # print(f"context is {self.tokenizer.decode(context_token_ids)}")
+            policy_output = await self.inference_backend.generate(
+                input_token_ids=context_token_ids.tolist(),
+                extract_thinking=(self.max_thinking_characters > 0),
+                regex=current_regex,
+            )
+            # print(f"generated: {self.tokenizer.decode(policy_output.out_token_ids)}")
+            if (
+                pattern is None
+                or (pattern.fullmatch(policy_output.content))
+                or (nb_attempts >= self.regex_max_attempts)
+            ):
+                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
+
+        # import random
+        # self.save_path = self.save_path + str(random.randint(1,500))
+        # print(f"Save path: {self.save_path}")
+        # self.adapter_paths = {adapter_id:os.path.join(self.save_path, adapter_id) for adapter_id in self.adapter_ids}
+
+        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,54 @@
+import torch, torch.nn as nn, torch.optim as optim
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import LoraConfig, get_peft_model
+
+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/README.md

@@ -0,0 +1,20 @@
+Suppose we have a trajectory with 3 timesteps.
+token:          "0 1 2 3 4 5 6 7 8 9 . . . . ."
+string:         "A B C a b c A a A a b c A B C" (Capitalized = User, Lowercased = Assistant)
+action_mask:    "x x x ✓ ✓ ✓ x ✓ x ✓ ✓ ✓ x x x" (F = False, T = True)
+rewards:        "r r r r r r R R R R R R r r r"
+timestep:       "0 0 0 0 0 0 1 1 1 1 1 1 2 2 2"
+state_ends:     "x x ✓ x x x ✓ x x x x x x x ✓"
+
+There must be one baseline flag per timestep!
+
+Then, we might have
+
+A naive way to interpret this is to think of the number of assistant messages as the number of
+steps in the environment. However, this is not the case in practice. Indeed, in a
+single simulation step,
+
+
+
+
+A subtlety arises with credit assignment. In the multi-agent case, we might