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: 10
+  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: tas_rps_startend_ad_align_nocurrtimestep_beta2
+  agent_buffer: true
+  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/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
+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
+      enable_thinking: null
+      regex_max_attempts: 1
+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: TrainerAdAlign
+    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
+      ad_align_force_coop_first_step: false
+      ad_align_clipping: null
+      ad_align_gamma: 0.96
+      ad_align_exclude_k_equals_t: true
+      ad_align_use_sign: false
+      ad_align_beta: 2.0
+      use_old_ad_align: true
+      use_time_regularization: false
+      rloo_branch: false
+      reuse_baseline: false
+train_on_which_data:
+  agent_trainer: ${agent_ids}
+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

.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: tas_rps_startend_ad_align_nocurrtimestep_beta2.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/tas_rps_startend_ad_align_nocurrtimestep_beta2
+    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

.hydra/overrides.yaml

@@ -0,0 +1 @@
+[]

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": [
+    "up_proj",
+    "q_proj",
+    "v_proj",
+    "gate_proj",
+    "down_proj",
+    "k_proj",
+    "o_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": [
+    "up_proj",
+    "q_proj",
+    "v_proj",
+    "gate_proj",
+    "down_proj",
+    "k_proj",
+    "o_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/docs/source/environments/diplomacy.rst

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

src_code_for_reproducibility/docs/source/environments/dond.rst

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

src_code_for_reproducibility/docs/source/environments/ipd.rst

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

src_code_for_reproducibility/docs/source/modules.rst

@@ -0,0 +1,7 @@
+src
+===
+
+.. toctree::
+   :maxdepth: 4
+
+   src

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

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

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

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

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

@@ -0,0 +1,7 @@
+src.models.dummy\_local\_llm module
+===================================
+
+.. automodule:: src.models.dummy_local_llm
+   :members:
+   :undoc-members:
+   :show-inheritance:

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

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

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

@@ -0,0 +1,7 @@
+src.models.local\_llm module
+============================
+
+.. automodule:: src.models.local_llm
+   :members:
+   :undoc-members:
+   :show-inheritance:

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

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

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

@@ -0,0 +1,7 @@
+src.models.server\_llm module
+=============================
+
+.. automodule:: src.models.server_llm
+   :members:
+   :undoc-members:
+   :show-inheritance:

src_code_for_reproducibility/docs/source/src.run.rst

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

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

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

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

@@ -0,0 +1,7 @@
+src.training.train\_main module
+===============================
+
+.. automodule:: src.training.train_main
+   :members:
+   :undoc-members:
+   :show-inheritance:

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

@@ -0,0 +1,7 @@
+src.utils.common\_imports module
+================================
+
+.. automodule:: src.utils.common_imports
+   :members:
+   :undoc-members:
+   :show-inheritance:

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,7 @@
+src.utils.model\_to\_cpu module
+===============================
+
+.. automodule:: src.utils.model_to_cpu
+   :members:
+   :undoc-members:
+   :show-inheritance:

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

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

src_code_for_reproducibility/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/training/tally_tokenwise.py

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

src_code_for_reproducibility/training/tokenize_chats.py

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

src_code_for_reproducibility/training/trainer_ad_align.py

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

src_code_for_reproducibility/training/trainer_common.py

@@ -0,0 +1,1054 @@
+"""
+TODO: Add coefficients for losses (depend on total number of tokens or batch)
+TODO: adapt reinforce step for torch.compile
+TODO: add lr schedulers support
+"""
+import logging
+import os
+import pickle
+import sys
+from abc import ABC, abstractmethod
+from typing import Callable, Literal, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from accelerate import Accelerator
+from pandas._libs.tslibs.offsets import CBMonthBegin
+from peft import LoraConfig
+from torch.nn.utils.rnn import pad_sequence
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from mllm.markov_games.rollout_tree import *
+from mllm.markov_games.rollout_tree import RolloutTreeRootNode
+from mllm.training.annealing_methods import sigmoid_annealing
+from mllm.training.credit_methods import (
+    get_discounted_returns,
+    get_generalized_advantage_estimates,
+    get_rloo_credits,
+    whiten_advantages,
+    whiten_advantages_time_step_wise,
+)
+from mllm.training.tally_metrics import Tally
+from mllm.training.tally_rollout import RolloutTally, RolloutTallyItem
+from mllm.training.tally_tokenwise import ContextualizedTokenwiseTally
+from mllm.training.tokenize_chats import *
+from mllm.training.tokenize_chats import process_training_chat
+from mllm.training.training_data_utils import *
+from mllm.training.training_data_utils import (
+    TrainingBatch,
+    TrajectoryBatch,
+    get_tokenwise_credits,
+)
+from mllm.utils.resource_context import resource_logger_context
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+@dataclass
+class TrainerAnnealingState:
+    annealing_step_counter: int = 0
+
+
+class BaseTrainer(ABC):
+    """
+    Trainer
+    """
+
+    def __init__(
+        self,
+        policy: AutoModelForCausalLM,
+        policy_optimizer: torch.optim.Optimizer,
+        critic: Union[AutoModelForCausalLM, None],
+        critic_optimizer: Union[torch.optim.Optimizer, None],
+        tokenizer: AutoTokenizer,
+        lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
+        critic_lr_scheduler: Union[torch.optim.lr_scheduler.LRScheduler, None],
+        ######################################################################
+        entropy_coeff: float,
+        entropy_topk: int,
+        entropy_mask_regex: Union[str, None],
+        kl_coeff: float,
+        gradient_clipping: Union[float, None],
+        restrict_tokens: Union[list[str], None],
+        mini_batch_size: int,
+        use_gradient_checkpointing: bool,
+        temperature: float,
+        device: str,
+        whiten_advantages: bool,
+        whiten_advantages_time_step_wise: bool,
+        use_gae: bool,
+        use_gae_lambda_annealing: bool,
+        gae_lambda_annealing_limit: float,
+        gae_lambda_annealing_method: Literal["sigmoid_annealing"],
+        gae_lambda_annealing_method_params: dict,
+        pg_loss_normalization: Literal["batch", "nb_tokens"],
+        use_rloo: bool,
+        skip_discounted_state_visitation: bool,
+        discount_factor: float,
+        enable_tokenwise_logging: bool,
+        save_path: str,
+        reward_normalizing_constant: float = 1.0,
+        critic_loss_type: Literal["mse", "huber"] = "huber",
+        exploration_prompts_to_remove: list[str] = [],
+        filter_higher_refprob_tokens_kl: bool = False,
+        truncated_importance_sampling_ratio_cap: float = 0.0,
+        importance_sampling_strategy: Literal[
+            "per_token", "per_sequence"
+        ] = "per_token",
+    ):
+        """
+        Initialize the REINFORCE trainer with reward shaping for multi-agent or single-agent training.
+
+        Args:
+            model (AutoModelForCausalLM): The main policy model.
+            tokenizer (AutoTokenizer): Tokenizer for the model.
+            optimizer (torch.optim.Optimizer): Optimizer for the policy model.
+            lr_scheduler (torch.optim.lr_scheduler.LRScheduler): Learning rate scheduler for the policy model.
+            critic (AutoModelForCausalLM or None): Critic model for value estimation (optional).
+            critic_optimizer (torch.optim.Optimizer or None): Optimizer for the critic model (optional).
+            critic_lr_scheduler (torch.optim.lr_scheduler.LRScheduler or None): LR scheduler for the critic (optional).
+            config (RtConfig): Configuration object for training.
+        """
+        self.tokenizer = tokenizer
+        # self.tokenizer.padding_side = "left"  # needed for flash attention
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.lr_scheduler = lr_scheduler
+        self.accelerator = Accelerator()
+        (
+            self.policy,
+            self.policy_optimizer,
+            self.critic,
+            self.critic_optimizer,
+        ) = self.accelerator.prepare(policy, policy_optimizer, critic, critic_optimizer)
+
+        self.critic_lr_scheduler = critic_lr_scheduler
+        self.tally = Tally()
+
+        if use_gradient_checkpointing == True:
+            self.policy.gradient_checkpointing_enable(dict(use_reentrant=False))
+            if critic is not None:
+                self.critic.gradient_checkpointing_enable(dict(use_reentrant=False))
+
+        self.save_path = save_path
+
+        # Load trainer state if it exists
+        self.trainer_annealing_state_path = os.path.join(
+            self.save_path, "trainer_annealing_state.pkl"
+        )
+        if os.path.exists(self.trainer_annealing_state_path):
+            logger.info(
+                f"Loading trainer state from {self.trainer_annealing_state_path}"
+            )
+            self.trainer_annealing_state = pickle.load(
+                open(self.trainer_annealing_state_path, "rb")
+            )
+        else:
+            self.trainer_annealing_state = TrainerAnnealingState()
+
+        # Load policy optimizer state if it exists
+        self.policy_optimizer_path = os.path.join(
+            self.save_path, "policy_optimizer_state.pt"
+        )
+        if os.path.exists(self.policy_optimizer_path):
+            logger.info(
+                f"Loading policy optimizer state from {self.policy_optimizer_path}"
+            )
+            self.policy_optimizer.load_state_dict(
+                torch.load(self.policy_optimizer_path)
+            )
+
+        # Load critic optimizer state if it exists
+        self.critic_optimizer_path = os.path.join(
+            self.save_path, "critic_optimizer_state.pt"
+        )
+        if (
+            os.path.exists(self.critic_optimizer_path)
+            and self.critic_optimizer is not None
+        ):
+            logger.info(
+                f"Loading critic optimizer state from {self.critic_optimizer_path}"
+            )
+            self.critic_optimizer.load_state_dict(
+                torch.load(self.critic_optimizer_path)
+            )
+        self.device = self.accelerator.device
+        self.entropy_coeff = entropy_coeff
+        self.entropy_topk = entropy_topk
+        self.entropy_mask_regex = entropy_mask_regex
+        self.kl_coeff = kl_coeff
+        self.gradient_clipping = gradient_clipping
+        self.restrict_tokens = restrict_tokens
+        self.mini_batch_size = mini_batch_size
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.temperature = temperature
+        self.use_gae = use_gae
+        self.whiten_advantages = whiten_advantages
+        self.whiten_advantages_time_step_wise = whiten_advantages_time_step_wise
+        self.use_rloo = use_rloo
+        self.skip_discounted_state_visitation = skip_discounted_state_visitation
+        self.use_gae_lambda_annealing = use_gae_lambda_annealing
+        self.gae_lambda_annealing_limit = gae_lambda_annealing_limit
+        if use_gae_lambda_annealing:
+            self.gae_lambda_annealing_method: Callable[
+                [int], float
+            ] = lambda step: eval(gae_lambda_annealing_method)(
+                step=step, **gae_lambda_annealing_method_params
+            )
+        self.discount_factor = discount_factor
+        self.enable_tokenwise_logging = enable_tokenwise_logging
+        self.reward_normalizing_constant = reward_normalizing_constant
+        self.pg_loss_normalization = pg_loss_normalization
+        self.critic_loss_type = critic_loss_type
+        self.exploration_prompts_to_remove = exploration_prompts_to_remove
+        # Common containers used by all trainers
+        self.training_data: dict = {}
+        self.debug_path_list: list[str] = []
+        self.policy_gradient_data = None
+        self.tally = Tally()
+        self.rollout_tally = RolloutTally()
+        self.tokenwise_tally: Union[ContextualizedTokenwiseTally, None] = None
+        self.filter_higher_refprob_tokens_kl = filter_higher_refprob_tokens_kl
+        self.truncated_importance_sampling_ratio_cap = (
+            truncated_importance_sampling_ratio_cap
+        )
+        self.importance_sampling_strategy = importance_sampling_strategy
+
+    def mask_non_restricted_token_logits(self, logits: torch.Tensor) -> torch.Tensor:
+        """
+        Masks logits so that only allowed tokens (as specified in config.restrict_tokens)
+        and the EOS token are active.
+        All other logits are set to -inf, effectively removing them from the softmax.
+
+        Args:
+            logits (torch.Tensor): The logits tensor of shape (B, S, V).
+
+        Returns:
+            torch.Tensor: The masked logits tensor.
+        """
+        # TODO: verify. Not sure what we do here is differentiable
+        # also, we recompute for nothing
+
+        if self.restrict_tokens is not None:
+            allowed_token_ids = []
+            for token in self.restrict_tokens:
+                token_ids = self.tokenizer(token, add_special_tokens=False)["input_ids"]
+                allowed_token_ids.append(token_ids[0])
+            allowed_token_ids.append(
+                self.tokenizer.eos_token_id
+            )  # This token should always be active
+            allowed_token_ids = torch.tensor(allowed_token_ids, device=logits.device)
+            # Mask log_probs and probs to only allowed tokens
+            mask = torch.zeros_like(logits).bool()  # (B, S, V)
+            mask[..., allowed_token_ids] = True
+            logits = torch.where(
+                mask,
+                logits,
+                torch.tensor(-float("inf"), device=logits.device),
+            )
+
+        return logits
+
+    # def get_gradient_magnitude(self, loss_term: torch.Tensor) -> float:
+    #     """
+    #     Computes the L2 norm of the gradients of the given loss term with respect to the model parameters.
+
+    #     Args:
+    #         loss_term (torch.Tensor): The loss tensor to compute gradients for.
+
+    #     Returns:
+    #         float: The L2 norm of the gradients, or 0.0 if no gradients are present.
+    #     """
+    #     with torch.no_grad():
+    #         grads = torch.autograd.grad(
+    #             loss_term,
+    #             [p for p in self.policy.parameters() if p.requires_grad],
+    #             retain_graph=True,
+    #             allow_unused=True,
+    #         )
+    #         grads = [g for g in grads if g is not None]
+    #         if not grads:
+    #             return torch.tensor(0.0, device=loss_term.device)
+    #         return torch.norm(torch.stack([g.norm(2) for g in grads])).item()
+
+    def apply_reinforce_step(
+        self,
+        training_batch: TrainingBatch,
+    ) -> None:
+        """
+        Applies a single REINFORCE policy gradient step using the provided batch of rollouts.
+        Handles batching, loss computation (including entropy and KL regularization), gradient accumulation, and optimizer step.
+        Optionally logs various metrics and statistics.
+
+        Args:
+            paths (list[str]): List of game complete file paths for each rollout.
+            contexts (list[torch.Tensor]): List of context tensors for each rollout.
+            credits (list[torch.Tensor]): List of credit tensors (rewards/advantages) for each rollout.
+            action_masks (list[torch.Tensor]): List of action mask tensors for each rollout.
+        """
+        with resource_logger_context(logger, "Apply reinforce step"):
+            self.policy.train()
+            mb_size = self.mini_batch_size
+            nb_rollouts = len(training_batch)
+
+            # Initialize running mean logs
+            running_mean_logs = {
+                "rl_objective": 0.0,
+                "policy_gradient_loss": 0.0,
+                "policy_gradient_norm": 0.0,
+                "log_probs": 0.0,
+                "credits": 0.0,
+                "entropy": 0.0,
+                "engine_log_probs_diff_clampfrac": 0.0,
+                "tis_imp_ratio": 0.0,
+                "ref_log_probs_diff_clampfrac": 0.0,
+                "higher_refprob_frac": 0.0,
+                "tis_imp_ratio_clampfrac": 0.0,
+            }
+            if self.entropy_coeff != 0.0:
+                running_mean_logs["entropy"] = 0.0
+            if self.kl_coeff != 0.0:
+                running_mean_logs["kl_divergence"] = 0.0
+
+            # Get total number of tokens generated
+            total_tokens_generated = 0
+            for att_mask in training_batch.batch_action_mask:
+                total_tokens_generated += att_mask.sum()
+
+            # Obtain loss normalization
+            if self.pg_loss_normalization == "nb_tokens":
+                normalization_factor = total_tokens_generated
+            elif self.pg_loss_normalization == "batch":
+                normalization_factor = np.ceil(nb_rollouts / mb_size).astype(int)
+            else:
+                raise ValueError(
+                    f"Invalid pg_loss_normalization: {self.pg_loss_normalization}"
+                )
+
+            # Gradient accumulation for each mini-batch
+            for mb in range(0, nb_rollouts, mb_size):
+                logger.info(f"Processing mini-batch {mb} of {nb_rollouts}")
+                loss = 0.0
+                training_mb = training_batch[mb : mb + mb_size]
+                training_mb = training_mb.get_padded_tensors()
+                training_mb.to(self.device)
+                (
+                    tokens_mb,
+                    action_mask_mb,
+                    entropy_mask_mb,
+                    credits_mb,
+                    engine_log_probs_mb,
+                    timesteps_mb,
+                ) = (
+                    training_mb.batch_input_ids,
+                    training_mb.batch_action_mask,
+                    training_mb.batch_entropy_mask,
+                    training_mb.batch_credits,
+                    training_mb.batch_engine_log_probs,
+                    training_mb.batch_timesteps,
+                )
+
+                # Next token prediction
+                contexts_mb = tokens_mb[:, :-1]
+                shifted_contexts_mb = tokens_mb[:, 1:]
+                action_mask_mb = action_mask_mb[:, 1:]
+                entropy_mask_mb = entropy_mask_mb[:, 1:]
+                credits_mb = credits_mb[:, 1:]
+                engine_log_probs_mb = engine_log_probs_mb[:, 1:]
+                timesteps_mb = timesteps_mb[:, 1:]
+
+                if self.enable_tokenwise_logging:
+                    self.tokenwise_tally.set_action_mask(action_mask=action_mask_mb)
+                    self.tokenwise_tally.set_range(range=(mb, mb + mb_size))
+                    self.tokenwise_tally.add_contexts(contexts=contexts_mb)
+                    self.tokenwise_tally.add_data(
+                        metric_id="next_token",
+                        metrics=shifted_contexts_mb,
+                        to_tids=True,
+                    )
+                    self.tokenwise_tally.add_data(
+                        metric_id="entropy_mask",
+                        metrics=entropy_mask_mb,
+                    )
+
+                if self.enable_tokenwise_logging:
+                    self.tokenwise_tally.add_data(
+                        metric_id="next_token_credit", metrics=credits_mb
+                    )
+
+                # Forward pass + cast to FP-32 for higher prec.
+                # TODO: create attention mask if not relying on default (assume causal llm)
+                logits = self.policy(input_ids=contexts_mb)[0]  # (B, S, V)
+
+                # Mask non-restricted tokens
+                if self.restrict_tokens is not None:
+                    logits = self.mask_non_restricted_token_logits(logits)
+
+                logits /= self.temperature  # (B, S, V)
+
+                # Compute new log probabilities
+                log_probs = F.log_softmax(logits, dim=-1)  # (B, S, V)
+
+                # Get log probabilities of actions taken during rollouts
+                action_log_probs = log_probs.gather(
+                    dim=-1, index=shifted_contexts_mb.unsqueeze(-1)
+                ).squeeze(
+                    -1
+                )  # (B, S)
+                if self.pg_loss_normalization == "batch":
+                    den_running_mean = action_mask_mb.sum() * normalization_factor
+                else:
+                    den_running_mean = normalization_factor
+                running_mean_logs["log_probs"] += (
+                    action_log_probs * action_mask_mb
+                ).sum().item() / den_running_mean
+                running_mean_logs["credits"] += (
+                    credits_mb * action_mask_mb
+                ).sum().item() / den_running_mean
+
+                if self.enable_tokenwise_logging:
+                    self.tokenwise_tally.add_data(
+                        metric_id="next_token_log_prob",
+                        metrics=action_log_probs,
+                    )
+                    self.tokenwise_tally.add_data(
+                        metric_id="engine_next_token_log_prob",
+                        metrics=engine_log_probs_mb,
+                    )
+                    self.tokenwise_tally.add_data(
+                        metric_id="next_token_prob",
+                        metrics=torch.exp(action_log_probs),
+                    )
+                    top_k_indices = torch.topk(logits, k=5, dim=-1).indices
+                    self.tokenwise_tally.add_data(
+                        metric_id=f"top_{5}_tids",
+                        metrics=top_k_indices,
+                        to_tids=True,
+                    )
+                    self.tokenwise_tally.add_data(
+                        metric_id=f"top_{5}_probs",
+                        metrics=torch.exp(log_probs).gather(
+                            dim=-1, index=top_k_indices
+                        ),
+                    )
+
+                rewarded_action_log_probs = (
+                    action_mask_mb * credits_mb * action_log_probs
+                )
+                # (B, S)
+                INVALID_LOGPROB = 1.0
+                CLAMP_VALUE = 40.0
+                masked_action_log_probs = torch.masked_fill(
+                    action_log_probs, ~action_mask_mb, INVALID_LOGPROB
+                )
+                masked_engine_log_probs = torch.masked_fill(
+                    engine_log_probs_mb, ~action_mask_mb, INVALID_LOGPROB
+                )
+                with torch.no_grad():
+                    action_engine_log_probs_diff = (
+                        masked_action_log_probs - masked_engine_log_probs
+                    ).clamp(-CLAMP_VALUE, CLAMP_VALUE)
+                running_mean_logs["engine_log_probs_diff_clampfrac"] += (
+                    action_engine_log_probs_diff.abs()
+                    .eq(CLAMP_VALUE)
+                    .float()
+                    .sum()
+                    .item()
+                    / den_running_mean
+                )
+                if self.importance_sampling_strategy == "per_sequence":
+                    tis_imp_ratio = torch.zeros_like(action_engine_log_probs_diff)
+                    for mb_idx in range(action_engine_log_probs_diff.shape[0]):
+                        valid_token_mask = action_mask_mb[mb_idx]
+                        timestep_ids = timesteps_mb[mb_idx][valid_token_mask]
+                        timestep_logprob_diffs = action_engine_log_probs_diff[mb_idx][
+                            valid_token_mask
+                        ]
+                        max_timestep = int(timestep_ids.max().item()) + 1
+                        timestep_sums = torch.zeros(
+                            max_timestep,
+                            device=action_engine_log_probs_diff.device,
+                            dtype=action_engine_log_probs_diff.dtype,
+                        )
+                        timestep_sums.scatter_add_(
+                            0, timestep_ids, timestep_logprob_diffs
+                        )
+                        timestep_ratios = torch.exp(timestep_sums)
+                        tis_imp_ratio[
+                            mb_idx, valid_token_mask
+                        ] = timestep_ratios.gather(0, timestep_ids)
+                else:
+                    tis_imp_ratio = torch.exp(action_engine_log_probs_diff)
+                running_mean_logs["tis_imp_ratio"] += (
+                    tis_imp_ratio * action_mask_mb
+                ).sum().item() / den_running_mean
+                if self.truncated_importance_sampling_ratio_cap > 0.0:
+                    tis_imp_ratio = torch.clamp(
+                        tis_imp_ratio, max=self.truncated_importance_sampling_ratio_cap
+                    )
+                    running_mean_logs["tis_imp_ratio_clampfrac"] += (
+                        tis_imp_ratio.eq(self.truncated_importance_sampling_ratio_cap)
+                        .float()
+                        .sum()
+                        .item()
+                    ) / den_running_mean
+                    rewarded_action_log_probs = (
+                        rewarded_action_log_probs * tis_imp_ratio
+                    )
+
+                if self.enable_tokenwise_logging:
+                    self.tokenwise_tally.add_data(
+                        metric_id="next_token_clogπ",
+                        metrics=rewarded_action_log_probs,
+                    )
+
+                # Add value term to loss
+                if self.pg_loss_normalization == "batch":
+                    nb_act_tokens = action_mask_mb.sum()
+                    mb_value = -rewarded_action_log_probs.sum() / nb_act_tokens
+                else:
+                    mb_value = -rewarded_action_log_probs.sum()
+
+                loss += mb_value
+                running_mean_logs["rl_objective"] += mb_value.item() / den_running_mean
+
+                # -------------------------------------------------
+                # Entropy Regularization
+                # -------------------------------------------------
+                # Only apply entropy on distribution defined over most probable tokens
+                if self.entropy_topk is not None:
+                    top_k_indices = torch.topk(
+                        logits, k=self.entropy_topk, dim=-1
+                    ).indices
+                    entropy_logits = logits.gather(dim=-1, index=top_k_indices)
+                else:
+                    entropy_logits = logits
+
+                token_entropy_terms = -F.softmax(
+                    entropy_logits, dim=-1
+                ) * F.log_softmax(
+                    entropy_logits, dim=-1
+                )  # (B, S, T)
+                token_entropy_terms *= (
+                    action_mask_mb[:, :, None] * entropy_mask_mb[:, :, None]
+                )  # only get loss on specific action tokens
+
+                mb_entropy = token_entropy_terms.sum(dim=-1)
+
+                if self.enable_tokenwise_logging:
+                    self.tokenwise_tally.add_data(
+                        metric_id="entropy",
+                        metrics=mb_entropy,
+                    )
+                if self.pg_loss_normalization == "batch":
+                    nb_act_tokens = action_mask_mb.sum()
+                    mb_entropy = -mb_entropy.sum() / nb_act_tokens
+                else:
+                    mb_entropy = -mb_entropy.sum()
+                running_mean_logs["entropy"] += -mb_entropy.item() / den_running_mean
+                if self.entropy_coeff != 0.0:
+                    mb_entropy *= self.entropy_coeff
+                    loss += mb_entropy
+
+                # -------------------------------------------------
+                # KL-DIVERGENCE
+                # -------------------------------------------------
+                if self.kl_coeff != 0.0:
+                    ref_model_logits = self.policy.get_base_model_logits(contexts_mb)
+                    ref_model_logits = ref_model_logits / self.temperature
+                    # (B, S, V)
+                    ref_model_logits = self.mask_non_restricted_token_logits(
+                        logits=ref_model_logits
+                    )
+                    # (B, S, V)
+                    ref_model_log_probs = F.log_softmax(ref_model_logits, dim=-1)
+                    # (B, S, V)
+                    ref_model_action_log_probs = ref_model_log_probs.gather(
+                        dim=-1, index=shifted_contexts_mb.unsqueeze(-1)
+                    ).squeeze(
+                        -1
+                    )  # (B,S)
+                    # Approximating KL Divergence (see refs in docstring)
+                    # Ref 1: http://joschu.net/blog/kl-approx.html
+                    # Ref 2: https://github.dev/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L1332
+                    masked_ref_model_action_log_probs = torch.masked_fill(
+                        ref_model_action_log_probs, ~action_mask_mb, INVALID_LOGPROB
+                    )
+                    action_log_probs_diff = (
+                        masked_ref_model_action_log_probs - masked_action_log_probs
+                    ).clamp(-CLAMP_VALUE, CLAMP_VALUE)
+                    running_mean_logs["ref_log_probs_diff_clampfrac"] += (
+                        action_log_probs_diff.abs().eq(CLAMP_VALUE).float().sum().item()
+                        / den_running_mean
+                    )
+                    if self.filter_higher_refprob_tokens_kl:
+                        higher_refprob_tokens_mask = action_log_probs_diff > 0.0
+                        running_mean_logs["higher_refprob_frac"] += (
+                            higher_refprob_tokens_mask.sum().item() / den_running_mean
+                        )
+                        action_log_probs_diff = action_log_probs_diff * (
+                            ~higher_refprob_tokens_mask
+                        )
+                    kl_div = torch.expm1(action_log_probs_diff) - action_log_probs_diff
+                    kl_div *= action_mask_mb  # We only care about KLD of action tokens
+                    if self.truncated_importance_sampling_ratio_cap > 0.0:
+                        kl_div = kl_div * tis_imp_ratio
+                    kl_div *= self.kl_coeff
+                    if self.enable_tokenwise_logging:
+                        self.tokenwise_tally.add_data(
+                            metric_id="ref_model_next_token_log_prob",
+                            metrics=ref_model_action_log_probs,
+                        )
+                        self.tokenwise_tally.add_data(
+                            metric_id="kl_divergence",
+                            metrics=kl_div,
+                        )
+
+                    if self.pg_loss_normalization == "batch":
+                        nb_act_tokens = action_mask_mb.sum()
+                        mb_kl = kl_div.sum() / nb_act_tokens
+                    else:
+                        mb_kl = kl_div.sum()
+                    running_mean_logs["kl_divergence"] += (
+                        mb_kl.item() / den_running_mean
+                    )
+                    loss += mb_kl
+
+                # Accumulate gradient
+                running_mean_logs["policy_gradient_loss"] += (
+                    loss.item() / den_running_mean
+                )
+                loss /= normalization_factor
+                self.accelerator.backward(loss)
+
+                # ensure gpu memory is freed
+                del training_mb
+                del log_probs
+                del logits
+                del loss
+                del action_log_probs
+                del rewarded_action_log_probs
+
+            logger.info(
+                f"Accumulated the policy gradient loss for {total_tokens_generated} tokens."
+            )
+
+            # Clip gradients and take step
+            if self.gradient_clipping is not None:
+                grad_norm = self.accelerator.clip_grad_norm_(
+                    self.policy.parameters(), self.gradient_clipping
+                )
+                running_mean_logs["policy_gradient_norm"] += grad_norm.item()
+
+            # Take step
+            self.policy_optimizer.step()
+            self.policy_optimizer.zero_grad()
+
+            # Store logs
+            for key, value in running_mean_logs.items():
+                self.tally.add_metric(path=key, metric=value)
+
+            # Clear
+            # TODO: verify
+            self.accelerator.clear(self.policy, self.policy_optimizer)
+            import gc
+
+            gc.collect()
+            torch.cuda.empty_cache()
+            return running_mean_logs
+
+    def get_advantages_with_critic_gradient_accumulation(
+        self, trajectories: TrajectoryBatch, critic_loss_scaling_factor: float = 2.0
+    ) -> torch.FloatTensor:
+        """
+        TOWRITE
+        Uses GAE if enabled, otherwise uses Monte Carlo returns.
+        Optionally trains the critic if GAE is used.
+        Returns:
+            advantages: NestedFloatTensors
+        """
+
+        mb_size = self.mini_batch_size
+        batch_size = trajectories.rollout_ids.shape[0]
+        agent_id = trajectories.agent_ids[0]
+        batch_rewards = trajectories.batch_rewards
+
+        ######################################
+        # use critic for advantage estimation
+        ######################################
+        if self.use_gae:
+            if "buffer" in agent_id:
+                self.critic.eval()
+                training = False
+            else:
+                self.critic.train()
+                training = True
+            advantages = []
+            # critic_loss_scaling_factor comes learning single critic for two agents
+            normalization_factor = (
+                np.ceil(batch_size / mb_size).astype(int) * critic_loss_scaling_factor
+            )
+            # For each minibatch
+            for mb in range(0, batch_size, mb_size):
+                trajectory_mb = trajectories[mb : mb + mb_size]
+                trajectory_mb.to(self.device)
+                rewards_mb = trajectory_mb.batch_rewards
+                (
+                    tokens_mb,
+                    state_ends_mask_mb,
+                    timestep_counts,
+                ) = trajectory_mb.get_padded_tensors_for_critic()
+                # critic causal attention up to end flags
+                if training:
+                    vals_estimate_full = self.critic(tokens_mb)
+                else:
+                    with torch.no_grad():
+                        vals_estimate_full = self.critic(tokens_mb)
+
+                # if vals_estimate_full.dim() == 3:
+                #     vals_estimate_full = vals_estimate_full.squeeze(-1)
+
+                # Select only positions where states end, per sample → list of (jT,)
+                B = tokens_mb.shape[0]
+                vals_list = [
+                    vals_estimate_full[b][state_ends_mask_mb[b]] for b in range(B)
+                ]
+
+                # Pad to (B, max_jT) = (B, S)
+                vals_estimate_mb = pad_sequence(
+                    vals_list, batch_first=True, padding_value=0.0
+                )
+                dtype = vals_estimate_mb.dtype
+                rewards_mb = pad_sequence(
+                    rewards_mb, batch_first=True, padding_value=0.0
+                ).to(
+                    dtype=dtype
+                )  # (B, S)
+                self.rollout_tally.add_metric(
+                    path=["batch_rewards"],
+                    rollout_tally_item=RolloutTallyItem(
+                        crn_ids=trajectory_mb.crn_ids,
+                        rollout_ids=trajectory_mb.rollout_ids,
+                        agent_ids=trajectory_mb.agent_ids,
+                        metric_matrix=rewards_mb,
+                    ),
+                )
+                if self.reward_normalizing_constant != 1.0:
+                    rewards_mb /= self.reward_normalizing_constant
+
+                det_vals_estimate_mb = vals_estimate_mb.detach()  # (B, max_jT)
+                self.rollout_tally.add_metric(
+                    path=["mb_value_estimates_critic"],
+                    rollout_tally_item=RolloutTallyItem(
+                        crn_ids=trajectory_mb.crn_ids,
+                        rollout_ids=trajectory_mb.rollout_ids,
+                        agent_ids=trajectory_mb.agent_ids,
+                        metric_matrix=det_vals_estimate_mb,
+                    ),
+                )
+
+                # Append a 0 value to the end of the value estimates
+                if det_vals_estimate_mb.shape[1] == rewards_mb.shape[1]:
+                    Bsize = det_vals_estimate_mb.shape[0]
+                    device = det_vals_estimate_mb.device
+                    dtype = det_vals_estimate_mb.dtype
+                    det_vals_estimate_mb = torch.cat(
+                        [
+                            det_vals_estimate_mb,
+                            torch.zeros((Bsize, 1), device=device, dtype=dtype),
+                        ],
+                        dim=1,
+                    )  # (B, max_jT+1)
+                else:
+                    raise ValueError(
+                        "Incompatible shapes for value estimates and rewards."
+                    )
+
+                # Get annealed lambda
+                if self.use_gae_lambda_annealing:
+                    annealing_constant = self.gae_lambda_annealing_method(
+                        step=self.trainer_annealing_state.annealing_step_counter
+                    )
+                    annealed_lambda = (
+                        self.gae_lambda_annealing_limit * annealing_constant
+                    )
+                    self.tally.add_metric(
+                        path="annealed_lambda", metric=annealed_lambda
+                    )
+                else:
+                    annealed_lambda = self.gae_lambda_annealing_limit
+
+                # Get GAE advantages
+                gae_advantages = get_generalized_advantage_estimates(
+                    rewards=rewards_mb,
+                    value_estimates=det_vals_estimate_mb,
+                    discount_factor=self.discount_factor,
+                    lambda_coef=annealed_lambda,
+                )  # (B, max_jT)
+                self.rollout_tally.add_metric(
+                    path=["mb_gae_advantages"],
+                    rollout_tally_item=RolloutTallyItem(
+                        crn_ids=trajectory_mb.crn_ids,
+                        rollout_ids=trajectory_mb.rollout_ids,
+                        agent_ids=trajectory_mb.agent_ids,
+                        metric_matrix=gae_advantages,
+                    ),
+                )
+                if training:
+                    targets = (
+                        gae_advantages.to(dtype=dtype) + det_vals_estimate_mb[:, :-1]
+                    )  # (B, max_jT) # A(s, a, b) + V(s) = Q(s, a, b)
+                    self.rollout_tally.add_metric(
+                        path=["mb_targets_critic"],
+                        rollout_tally_item=RolloutTallyItem(
+                            crn_ids=trajectory_mb.crn_ids,
+                            rollout_ids=trajectory_mb.rollout_ids,
+                            agent_ids=trajectory_mb.agent_ids,
+                            metric_matrix=targets,
+                        ),
+                    )
+                    if self.critic_loss_type == "mse":
+                        loss = F.mse_loss(
+                            input=vals_estimate_mb,
+                            target=targets,
+                        )
+                    elif self.critic_loss_type == "huber":
+                        loss = F.huber_loss(
+                            input=vals_estimate_mb,
+                            target=targets,
+                        )
+                    self.tally.add_metric(path=["mb_critic_loss"], metric=loss.item())
+                    # Accumulate gradient
+                    loss /= normalization_factor
+                    self.accelerator.backward(loss)
+                    del loss
+                    del targets
+                    del vals_estimate_mb
+                del trajectory_mb
+                del vals_estimate_full
+
+                # Get jagged back using timestep_counts
+                advantages.extend(
+                    [gae_advantages[i, : timestep_counts[i]] for i in range(B)]
+                )
+
+        ######################################
+        # use exclusively Monte Carlo returns & rloo for advantage estimation
+        ######################################
+        else:
+            lengths = [len(c) for c in batch_rewards]
+            padded_rewards = pad_sequence(
+                batch_rewards, batch_first=True, padding_value=0.0
+            )
+            self.rollout_tally.add_metric(
+                path=["mb_rewards"],
+                rollout_tally_item=RolloutTallyItem(
+                    crn_ids=trajectories.crn_ids,
+                    rollout_ids=trajectories.rollout_ids,
+                    agent_ids=trajectories.agent_ids,
+                    metric_matrix=padded_rewards,
+                ),
+            )
+            if self.reward_normalizing_constant != 1.0:
+                padded_rewards /= self.reward_normalizing_constant
+            padded_advantages = get_discounted_returns(
+                rewards=padded_rewards,
+                discount_factor=self.discount_factor,
+            )  # no baseline for now
+            if self.use_rloo:
+                is_grouped_by_rng = (
+                    trajectories.crn_ids.unique().shape[0]
+                    != trajectories.crn_ids.shape[0]
+                )
+                if is_grouped_by_rng:
+                    for crn_id in trajectories.crn_ids.unique():
+                        rng_mask = trajectories.crn_ids == crn_id
+                        rng_advantages = padded_advantages[rng_mask]
+                        rng_advantages, _ = get_rloo_credits(credits=rng_advantages)
+                        padded_advantages[rng_mask] = rng_advantages
+                else:
+                    padded_advantages, _ = get_rloo_credits(credits=padded_advantages)
+                self.rollout_tally.add_metric(
+                    path=["mb_rloo_advantages"],
+                    rollout_tally_item=RolloutTallyItem(
+                        crn_ids=trajectories.crn_ids,
+                        rollout_ids=trajectories.rollout_ids,
+                        agent_ids=trajectories.agent_ids,
+                        metric_matrix=padded_advantages,
+                    ),
+                )
+            advantages = [
+                padded_advantages[i, : lengths[i]]
+                for i in range(padded_advantages.shape[0])
+            ]
+
+        if self.whiten_advantages_time_step_wise or self.whiten_advantages:
+            lengths = [len(c) for c in advantages]
+            padded_advantages = pad_sequence(
+                advantages, batch_first=True, padding_value=0.0
+            )
+            if self.whiten_advantages_time_step_wise:
+                whitened_padded_advantages = whiten_advantages_time_step_wise(
+                    padded_advantages
+                )
+                path = ["mb_whitened_advantages_time_step_wise"]
+            elif self.whiten_advantages:
+                whitened_padded_advantages = whiten_advantages(padded_advantages)
+                path = ["mb_whitened_advantages"]
+            self.rollout_tally.add_metric(
+                path=path,
+                rollout_tally_item=RolloutTallyItem(
+                    crn_ids=trajectories.crn_ids,
+                    rollout_ids=trajectories.rollout_ids,
+                    agent_ids=trajectories.agent_ids,
+                    metric_matrix=whitened_padded_advantages,
+                ),
+            )
+            advantages = [
+                whitened_padded_advantages[i, : lengths[i]]
+                for i in range(whitened_padded_advantages.shape[0])
+            ]
+
+        self.trainer_annealing_state.annealing_step_counter += 1
+
+        return advantages
+
+    @abstractmethod
+    def set_agent_trajectory_data(
+        self, agent_id: str, roots: list[RolloutTreeRootNode]
+    ) -> None:
+        """
+        TOWRITE
+        """
+        pass
+
+    def set_trajectory_data(
+        self, roots: list[RolloutTreeRootNode], agent_ids: list[str]
+    ) -> None:
+        """
+        TOWRITE
+        """
+        for agent_id in agent_ids:
+            self.set_agent_trajectory_data(agent_id, roots)
+
+    @abstractmethod
+    def share_advantage_data(self) -> list[AdvantagePacket]:
+        pass
+
+    @abstractmethod
+    def receive_advantage_data(self, advantage_packets: list[AdvantagePacket]) -> None:
+        pass
+
+    def set_policy_gradient_data(self, agent_ids: list[str]) -> None:
+        """
+        Already set earlier # TODO: make it separate and clean
+        """
+        self.policy_gradient_data = None
+        # for agent_id, trajectory_batch in self.training_data.items():
+        #     if "buffer" in agent_id:
+        #         continue
+        for agent_id in agent_ids:
+            assert "buffer" not in agent_id, "Buffer agents do not train policy"
+            trajectory_batch = self.training_data[agent_id]
+            tokenwise_batch_credits = get_tokenwise_credits(
+                batch_timesteps=trajectory_batch.batch_timesteps,
+                batch_credits=trajectory_batch.batch_credits,
+            )
+            policy_gradient_data = TrainingBatch(
+                rollout_ids=trajectory_batch.rollout_ids,
+                batch_input_ids=trajectory_batch.batch_input_ids,
+                batch_action_mask=trajectory_batch.batch_action_mask,
+                batch_entropy_mask=trajectory_batch.batch_entropy_mask,
+                batch_credits=tokenwise_batch_credits,
+                batch_engine_log_probs=trajectory_batch.batch_engine_log_probs,
+                batch_timesteps=trajectory_batch.batch_timesteps,
+            )
+            if self.policy_gradient_data is None:
+                self.policy_gradient_data = policy_gradient_data
+            else:
+                self.policy_gradient_data.append(policy_gradient_data)
+
+        self.training_data = {}
+        self.tokenwise_tally = ContextualizedTokenwiseTally(
+            tokenizer=self.tokenizer,
+            paths=self.debug_path_list,
+        )
+
+    def train(self) -> None:
+        """
+        TOWRITE
+        """
+        assert self.policy_gradient_data is not None, "Policy gradient data is not set"
+        if self.critic_optimizer is not None:
+            if self.gradient_clipping is not None:
+                grad_norm = self.accelerator.clip_grad_norm_(
+                    self.critic.parameters(), self.gradient_clipping
+                )
+                self.tally.add_metric(
+                    path="gradient_norm_critic", metric=grad_norm.item()
+                )
+            # Take step
+            self.critic_optimizer.step()
+            self.critic_optimizer.zero_grad()
+            self.accelerator.clear(self.critic, self.critic_optimizer)
+            import gc
+
+            gc.collect()
+            torch.cuda.empty_cache()
+        running_mean_logs = self.apply_reinforce_step(
+            training_batch=self.policy_gradient_data
+        )
+        return running_mean_logs
+
+    def export_training_tally(self, identifier: str, folder: str) -> None:
+        """
+        Saves and resets the collected training metrics using the tally object.
+        """
+        os.makedirs(folder, exist_ok=True)
+        self.tally.save(identifier=identifier, folder=folder)
+        self.tokenwise_tally.save(
+            path=os.path.join(folder, f"{identifier}_tokenwise.csv")
+        )
+        self.rollout_tally.save(identifier=identifier, folder=folder)
+        self.tally.reset()
+        self.tokenwise_tally = None
+        self.rollout_tally.reset()
+        self.debug_path_list = []
+
+    def export_optimizer_states(self) -> None:
+        """
+        Saves the optimizer states for both the main model and critic (if it exists).
+        """
+        try:
+            os.makedirs(self.save_path, exist_ok=True)
+
+            torch.save(self.policy_optimizer.state_dict(), self.policy_optimizer_path)
+            logger.info(f"Saved main optimizer state to {self.policy_optimizer_path}")
+
+            if self.critic_optimizer is not None:
+                torch.save(
+                    self.critic_optimizer.state_dict(), self.critic_optimizer_path
+                )
+                logger.info(
+                    f"Saved critic optimizer state to {self.critic_optimizer_path}"
+                )
+        except Exception as e:
+            logger.error(f"Error saving optimizer states: {str(e)}")
+            raise
+
+    def export_trainer_annealing_state(self) -> None:
+        """
+        Saves the trainer state.
+        """
+        with open(self.trainer_annealing_state_path, "wb") as f:
+            pickle.dump(self.trainer_annealing_state, f)
+        logger.info(f"Saved trainer state to {self.trainer_annealing_state_path}")
+
+    def export_trainer_states(self) -> None:
+        """
+        Saves the trainer states.
+        """
+        self.export_optimizer_states()
+        self.export_trainer_annealing_state()

src_code_for_reproducibility/training/trainer_independent.py

@@ -0,0 +1,155 @@
+"""
+
+"""
+import logging
+import os
+import sys
+from typing import Union
+
+import torch
+import torch.nn.functional as F
+from accelerate import Accelerator
+from pandas._libs.tslibs.offsets import CBMonthBegin
+from peft import LoraConfig
+from torch.nn.utils.rnn import pad_sequence
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from mllm.markov_games.rollout_tree import *
+from mllm.markov_games.rollout_tree import RolloutTreeRootNode
+from mllm.training.credit_methods import (
+    get_discounted_returns,
+    get_discounted_state_visitation_credits,
+    get_generalized_advantage_estimates,
+    get_rloo_credits,
+)
+from mllm.training.tally_metrics import Tally
+from mllm.training.tally_tokenwise import ContextualizedTokenwiseTally
+from mllm.training.tokenize_chats import *
+from mllm.training.tokenize_chats import process_training_chat
+from mllm.training.trainer_common import BaseTrainer
+from mllm.training.training_data_utils import *
+from mllm.training.training_data_utils import (
+    TrainingBatch,
+    TrajectoryBatch,
+    get_tokenwise_credits,
+)
+from mllm.utils.resource_context import resource_logger_context
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+@dataclass
+class TrainingData:
+    agent_id: str
+    main_data: TrajectoryBatch
+    # list-of-tensors: per rollout advantages with length jT
+    main_advantages: list[torch.FloatTensor] | None = None
+
+
+class TrainerNaive(BaseTrainer):
+    def set_agent_trajectory_data(
+        self, agent_id: str, roots: list[RolloutTreeRootNode]
+    ) -> None:
+        """
+        TOWRITE
+        """
+        # TODO: append to current batch data instead, else we will only train for one agent!
+        self.policy_gradient_data = None
+
+        # Tensorize Chats
+        rollout_ids = []
+        crn_ids = []  # common random number id
+        batch_input_ids = []
+        batch_action_mask = []
+        batch_entropy_mask = []
+        batch_timesteps = []
+        batch_state_ends_mask = []
+        batch_engine_log_probs = []
+        batch_rewards = []
+        for root in roots:
+            rollout_id = root.id
+            self.debug_path_list.append(
+                "mgid:" + str(rollout_id) + "_agent_id:" + agent_id
+            )
+            rollout_ids.append(rollout_id)
+            crn_ids.append(root.crn_id)
+            chat, rewards = get_main_chat_list_and_rewards(agent_id=agent_id, root=root)
+            (
+                input_ids,
+                action_mask,
+                entropy_mask,
+                timesteps,
+                state_ends_mask,
+                engine_log_probs,
+            ) = process_training_chat(
+                tokenizer=self.tokenizer,
+                chat_history=chat,
+                entropy_mask_regex=self.entropy_mask_regex,
+                exploration_prompts_to_remove=self.exploration_prompts_to_remove,
+            )
+            batch_input_ids.append(input_ids)
+            batch_action_mask.append(action_mask)
+            batch_entropy_mask.append(entropy_mask)
+            batch_timesteps.append(timesteps)
+            batch_state_ends_mask.append(state_ends_mask)
+            batch_engine_log_probs.append(engine_log_probs)
+            batch_rewards.append(rewards)
+
+        trajectory_batch = TrajectoryBatch(
+            rollout_ids=torch.tensor(rollout_ids, dtype=torch.int32),
+            crn_ids=torch.tensor(crn_ids, dtype=torch.int32),
+            agent_ids=[agent_id] * len(rollout_ids),
+            batch_input_ids=batch_input_ids,
+            batch_action_mask=batch_action_mask,
+            batch_entropy_mask=batch_entropy_mask,
+            batch_timesteps=batch_timesteps,
+            batch_state_ends_mask=batch_state_ends_mask,
+            batch_rewards=batch_rewards,
+            batch_engine_log_probs=batch_engine_log_probs,
+        )
+
+        # Get Advantages
+        batch_advantages: torch.FloatTensor = (
+            self.get_advantages_with_critic_gradient_accumulation(trajectory_batch)
+        )
+
+        # Discount state visitation (the mathematically correct way)
+        if not self.skip_discounted_state_visitation:
+            for i in range(len(batch_advantages)):
+                batch_advantages[i] = get_discounted_state_visitation_credits(
+                    batch_advantages[i].unsqueeze(0),
+                    self.discount_factor,
+                ).squeeze(0)
+
+        self.training_data[agent_id] = TrainingData(
+            agent_id=agent_id,
+            main_data=trajectory_batch,
+            main_advantages=batch_advantages,
+        )
+
+    def receive_advantage_data(self, advantage_packets: list[AdvantagePacket]):
+        """
+        This trainer ignores the advantages of the other trainers.
+        """
+        for agent_id, agent_data in self.training_data.items():
+            self.training_data[agent_id] = agent_data.main_data
+            self.training_data[agent_id].batch_credits = agent_data.main_advantages
+
+    def share_advantage_data(self) -> list[AdvantagePacket]:
+        """
+        Share the advantage data with other agents.
+        Returns:
+            AdvantagePacket: The advantage packet containing the agent's advantages.
+        """
+        logger.info(f"Sharing advantage data.")
+        advantage_packets = []
+        for agent_id, agent_data in self.training_data.items():
+            advantage_packets.append(
+                AdvantagePacket(
+                    agent_id=agent_id,
+                    rollout_ids=agent_data.main_data.rollout_ids,
+                    main_advantages=agent_data.main_advantages,
+                )
+            )
+        return advantage_packets

src_code_for_reproducibility/training/training_data_utils.py

@@ -0,0 +1,394 @@
+from dataclasses import dataclass
+from typing import Literal, Optional, Tuple
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from mllm.markov_games.rollout_tree import (
+    ChatTurn,
+    RolloutTreeBranchNode,
+    RolloutTreeNode,
+    RolloutTreeRootNode,
+)
+
+
+@dataclass
+class AdvantagePacket:
+    agent_id: str
+    rollout_ids: torch.IntTensor  # (B,)
+    # list-of-tensors
+    main_advantages: list[torch.FloatTensor]
+
+
+class TrainingChatTurn:
+    # TODO: simplify by making this a child of ChatTurn
+    """
+    This class contains the chat turns for a single agent.
+    It is like ChatTurn, but with the time step added.
+    """
+
+    def __init__(
+        self,
+        time_step: int,
+        role: str,
+        agent_id: str,
+        content: str,
+        chat_template_token_ids: list[int],
+        reasoning_content: str,
+        is_state_end: bool,
+        out_token_ids: Optional[list[int]] = None,
+        log_probs: Optional[list[float]] = None,
+    ) -> None:
+        self.time_step = time_step
+        self.role = role
+        self.agent_id = agent_id
+        self.content = content
+        self.chat_template_token_ids = chat_template_token_ids
+        self.reasoning_content = reasoning_content
+        self.is_state_end = is_state_end
+        self.out_token_ids = out_token_ids
+        self.log_probs = log_probs
+
+    def dict(self):
+        return {
+            "time_step": self.time_step,
+            "role": self.role,
+            "agent_id": self.agent_id,
+            "content": self.content,
+            "chat_template_token_ids": self.chat_template_token_ids,
+            "reasoning_content": self.reasoning_content,
+            "is_state_end": self.is_state_end,
+            "out_token_ids": self.out_token_ids,
+            "log_probs": self.log_probs,
+        }
+
+
+def get_main_chat_list_and_rewards(
+    agent_id: str, root: RolloutTreeRootNode | RolloutTreeNode
+) -> Tuple[list[TrainingChatTurn], torch.FloatTensor]:
+    """
+    This method traverses a rollout tree and returns a the list of ChatTurn
+    for an agent. If it encounters a branch node, it follows the main path.
+    """
+    # TODO; extend for all trees, not just linear
+    if isinstance(root, RolloutTreeRootNode):
+        current_node = root.child
+    else:
+        current_node = root
+
+    chat = []
+    rewards = []
+    while current_node is not None:
+        if isinstance(current_node, RolloutTreeBranchNode):
+            current_node = current_node.main_child
+        reward: float = current_node.step_log.simulation_step_log.rewards[agent_id]
+        rewards.append(reward)
+        chat_turns: list[TrainingChatTurn] = current_node.step_log.action_logs[
+            agent_id
+        ].chat_turns
+        chat_turns = [
+            TrainingChatTurn(time_step=current_node.time_step, **turn.model_dump())
+            for turn in chat_turns
+        ]
+        chat.extend(chat_turns)
+        current_node = current_node.child
+    return chat, torch.FloatTensor(rewards)
+
+
+def get_tokenwise_credits(
+    # B := batch size, S := number of tokens / seq. length, T := number of states. `j` stands for jagged (see pytorch nested tensors.)
+    batch_timesteps: torch.IntTensor | torch.Tensor,  # (B, jS),
+    batch_credits: torch.FloatTensor | torch.Tensor,  # (B, jT)
+) -> torch.FloatTensor | torch.Tensor:  # (B, jS)
+    """
+    TOWRITE
+    """
+    # TODO vectorize this code
+    batch_token_credits = []
+    for credits, timesteps in zip(batch_credits, batch_timesteps):
+        token_credits = torch.zeros_like(
+            timesteps,
+            dtype=credits.dtype,
+            device=timesteps.device,
+        )
+        for idx, credit in enumerate(credits):
+            token_credits[timesteps == idx] = credit
+        batch_token_credits.append(token_credits)
+    return batch_token_credits
+
+
+@dataclass
+class TrajectoryBatch:
+    """
+    Tensorized batch of trajectories using list-of-tensors for jagged dimensions.
+    """
+
+    # B := batch size, S := number of tokens / seq. length, T := number of states.
+    rollout_ids: torch.IntTensor  # (B,)
+    crn_ids: torch.IntTensor  # (B,)
+    agent_ids: list[str]  # (B,)
+    batch_input_ids: list[torch.LongTensor]  # List[(jS,)]
+    batch_action_mask: list[torch.BoolTensor]  # List[(jS,)]
+    batch_entropy_mask: list[torch.BoolTensor]  # List[(jS,)]
+    batch_timesteps: list[torch.IntTensor]  # List[(jS,)]
+    batch_state_ends_mask: list[torch.BoolTensor]  # List[(jS,)]
+    batch_engine_log_probs: Optional[list[torch.FloatTensor]]  # List[(jS,)]
+    batch_rewards: list[torch.FloatTensor]  # List[(jT,)]
+    batch_credits: Optional[list[torch.FloatTensor]] = None  # List[(jS,)]
+
+    def __post_init__(self):
+        """
+        Validate per-sample consistency.
+        """
+        B = self.rollout_ids.shape[0]
+        assert (
+            self.crn_ids.shape[0] == B
+        ), "RNG IDs must have length equal to batch size."
+        assert (
+            len(self.agent_ids) == B
+        ), "agent_ids must have length equal to batch size."
+        assert (
+            len(self.batch_input_ids)
+            == len(self.batch_action_mask)
+            == len(self.batch_entropy_mask)
+            == len(self.batch_timesteps)
+            == len(self.batch_state_ends_mask)
+            == len(self.batch_engine_log_probs)
+            == len(self.batch_rewards)
+            == B
+        ), "Jagged lists must all have length equal to batch size."
+
+        for b in range(B):
+            nb_rewards = int(self.batch_rewards[b].shape[0])
+            nb_timesteps = int(torch.max(self.batch_timesteps[b]).item()) + 1
+            assert (
+                nb_rewards == nb_timesteps
+            ), "Number of rewards and timesteps mismatch."
+            assert (
+                self.batch_input_ids[b].shape[0]
+                == self.batch_action_mask[b].shape[0]
+                == self.batch_entropy_mask[b].shape[0]
+                == self.batch_engine_log_probs[b].shape[0]
+                == self.batch_timesteps[b].shape[0]
+            ), "Tensors must have the same shape along the jagged dimension."
+            assert (
+                int(self.batch_state_ends_mask[b].sum())
+                == self.batch_rewards[b].shape[0]
+            ), "Number of rewards must match number of state ends."
+
+    """
+    Entries:
+        Here, we ignore the batch dimension.
+        input_ids:
+            All of the tokens of both the user and the assistant, flattened.
+        action_mask:
+            Set to true on the tokens of the assistant (tokens generated by the model).
+        timesteps:
+            Therefore, max(timesteps) = Ns - 1.
+        state_ends_idx:
+            Indices of the tokens at which state descriptions end.
+        rewards:
+            rewards[t] := R_t(s_t, a_t)
+    Example:
+        position:       "0  1  2  3  4  5  6  7  8  9  10 11 12 13 14"
+        input_ids:      "U  U  U  a  a  a  U  a  U  a  a  a  U  U  U" (U := User, a := Assistant)
+        action_mask:    "x  x  x  ✓  ✓  ✓  x  ✓  x  ✓  ✓  ✓  x  x  x"
+        timestep:       "0  0  0  0  0  0  1  1  1  1  1  1  2  2  2"
+        state_ends_dx:  [2, 6, 14]
+        rewards:        [r0, r1, r2]
+    """
+
+    def __getitem__(self, key) -> "TrajectoryBatch":
+        if isinstance(key, slice):
+            return TrajectoryBatch(
+                rollout_ids=self.rollout_ids.__getitem__(key),
+                crn_ids=self.crn_ids.__getitem__(key),
+                agent_ids=self.agent_ids[key],
+                batch_input_ids=self.batch_input_ids[key],
+                batch_action_mask=self.batch_action_mask[key],
+                batch_entropy_mask=self.batch_entropy_mask[key],
+                batch_timesteps=self.batch_timesteps[key],
+                batch_state_ends_mask=self.batch_state_ends_mask[key],
+                batch_engine_log_probs=self.batch_engine_log_probs[key],
+                batch_rewards=self.batch_rewards[key],
+                batch_credits=self.batch_credits[key] if self.batch_credits else None,
+            )
+
+    def __len__(self):
+        return len(self.batch_input_ids)
+
+    def to(self, device):
+        self.rollout_ids = self.rollout_ids.to(device)
+        self.crn_ids = self.crn_ids.to(device)
+        self.batch_input_ids = [t.to(device) for t in self.batch_input_ids]
+        self.batch_action_mask = [t.to(device) for t in self.batch_action_mask]
+        self.batch_entropy_mask = [t.to(device) for t in self.batch_entropy_mask]
+        self.batch_timesteps = [t.to(device) for t in self.batch_timesteps]
+        self.batch_state_ends_mask = [t.to(device) for t in self.batch_state_ends_mask]
+        self.batch_engine_log_probs = [
+            t.to(device) for t in self.batch_engine_log_probs
+        ]
+        self.batch_rewards = [t.to(device) for t in self.batch_rewards]
+        self.batch_credits = (
+            [t.to(device) for t in self.batch_credits] if self.batch_credits else None
+        )
+
+    def get_padded_tensors_for_critic(self):
+        """
+        Returns:
+            padded_batch_input_ids: (B, P)
+            padded_batch_state_ends_mask: (B, P)
+            timestep_counts: (B,) tensor of ints indicating number of states per sample
+        """
+        padded_batch_input_ids = pad_sequence(
+            self.batch_input_ids, batch_first=True, padding_value=0
+        )
+        padded_batch_state_ends_mask = pad_sequence(
+            self.batch_state_ends_mask, batch_first=True, padding_value=0
+        ).bool()
+        # number of states equals number of True in state_ends_mask
+        timestep_counts = torch.tensor(
+            [int(mask.sum().item()) for mask in self.batch_state_ends_mask],
+            device=padded_batch_input_ids.device,
+            dtype=torch.long,
+        )
+        return padded_batch_input_ids, padded_batch_state_ends_mask, timestep_counts
+
+
+timestep = int
+
+
+@dataclass
+class PaddedTensorTrainingBatch:
+    batch_input_ids: torch.LongTensor | torch.Tensor
+    batch_action_mask: torch.BoolTensor | torch.Tensor
+    batch_entropy_mask: Optional[torch.BoolTensor | torch.Tensor]
+    batch_credits: torch.FloatTensor | torch.Tensor
+    batch_engine_log_probs: torch.FloatTensor | torch.Tensor
+    batch_timesteps: torch.IntTensor | torch.Tensor
+
+    def __len__(self):
+        return self.batch_input_ids.shape[0]
+
+    def to(self, device):
+        self.batch_input_ids = self.batch_input_ids.to(device)
+        self.batch_action_mask = self.batch_action_mask.to(device)
+        self.batch_entropy_mask = self.batch_entropy_mask.to(device)
+        self.batch_credits = self.batch_credits.to(device)
+        self.batch_engine_log_probs = self.batch_engine_log_probs.to(device)
+        self.batch_timesteps = self.batch_timesteps.to(device)
+
+
+@dataclass
+class TrainingBatch:
+    rollout_ids: torch.IntTensor | torch.Tensor  # (B,)
+    batch_input_ids: list[torch.LongTensor]  # List[(jS,)]
+    batch_action_mask: list[torch.BoolTensor]  # List[(jS,)]
+    batch_entropy_mask: Optional[list[torch.BoolTensor]]  # List[(jS,)]
+    batch_credits: list[torch.FloatTensor]  # List[(jS,)]
+    batch_engine_log_probs: list[torch.FloatTensor]  # List[(jS,)]
+    batch_timesteps: list[torch.IntTensor]  # List[(jS,)]
+
+    def __post_init__(self):
+        # Put everything in the right device
+        # self.rollout_ids = self.rollout_ids.to("cuda" if torch.cuda.is_available() else "cpu")
+        # self.batch_input_ids = self.batch_input_ids.to("cuda" if torch.cuda.is_available() else "cpu")
+        # self.batch_action_mask = self.batch_action_mask.to("cuda" if torch.cuda.is_available() else "cpu")
+        # self.batch_credits = self.batch_credits.to("cuda" if torch.cuda.is_available() else "cpu")
+        # Ensure batch dimension is present
+        assert (
+            len(self.batch_input_ids)
+            == len(self.batch_action_mask)
+            == len(self.batch_entropy_mask)
+            == len(self.batch_credits)
+            == len(self.batch_engine_log_probs)
+            == len(self.batch_timesteps)
+            == self.rollout_ids.shape[0]
+        ), "Jagged lists must all have length equal to batch size."
+        for inp, mask, cred, engine_log_prob, timestep in zip(
+            self.batch_input_ids,
+            self.batch_action_mask,
+            self.batch_credits,
+            self.batch_engine_log_probs,
+            self.batch_timesteps,
+        ):
+            assert (
+                inp.shape[0]
+                == mask.shape[0]
+                == cred.shape[0]
+                == engine_log_prob.shape[0]
+                == timestep.shape[0]
+            ), "Tensors must have the same shapes along the jagged dimension."
+
+    def __getitem__(self, key) -> "TrainingBatch":
+        if isinstance(key, slice):
+            return TrainingBatch(
+                rollout_ids=self.rollout_ids.__getitem__(key),
+                batch_input_ids=self.batch_input_ids[key],
+                batch_action_mask=self.batch_action_mask[key],
+                batch_entropy_mask=self.batch_entropy_mask[key],
+                batch_credits=self.batch_credits[key],
+                batch_engine_log_probs=self.batch_engine_log_probs[key],
+                batch_timesteps=self.batch_timesteps[key],
+            )
+
+    def __len__(self):
+        return len(self.batch_input_ids)
+
+    def to(self, device):
+        self.rollout_ids = self.rollout_ids.to(device)
+        self.batch_input_ids = [t.to(device) for t in self.batch_input_ids]
+        self.batch_action_mask = [t.to(device) for t in self.batch_action_mask]
+        self.batch_entropy_mask = [t.to(device) for t in self.batch_entropy_mask]
+        self.batch_credits = [t.to(device) for t in self.batch_credits]
+        self.batch_engine_log_probs = [
+            t.to(device) for t in self.batch_engine_log_probs
+        ]
+        self.batch_timesteps = [t.to(device) for t in self.batch_timesteps]
+
+    def get_padded_tensors(self, padding: float = 0.0):
+        """
+        TOWRITE
+        Always pad to the right.
+        """
+        padded_batch_input_ids = pad_sequence(
+            self.batch_input_ids, batch_first=True, padding_value=int(padding)
+        )
+        padded_batch_action_mask = pad_sequence(
+            [m.to(dtype=torch.bool) for m in self.batch_action_mask],
+            batch_first=True,
+            padding_value=False,
+        )
+        padded_batch_entropy_mask = pad_sequence(
+            self.batch_entropy_mask, batch_first=True, padding_value=False
+        )
+        padded_batch_credits = pad_sequence(
+            self.batch_credits, batch_first=True, padding_value=float(padding)
+        )
+        padded_batch_engine_log_probs = pad_sequence(
+            self.batch_engine_log_probs, batch_first=True, padding_value=float(padding)
+        )
+        padded_batch_timesteps = pad_sequence(
+            self.batch_timesteps, batch_first=True, padding_value=0
+        )
+
+        return PaddedTensorTrainingBatch(
+            padded_batch_input_ids,
+            padded_batch_action_mask,
+            padded_batch_entropy_mask,
+            padded_batch_credits,
+            padded_batch_engine_log_probs,
+            padded_batch_timesteps,
+        )
+
+    def append(self, other: "TrainingBatch"):
+        self.rollout_ids = torch.cat([self.rollout_ids, other.rollout_ids])
+        self.batch_input_ids.extend(other.batch_input_ids)
+        self.batch_action_mask.extend(other.batch_action_mask)
+        self.batch_entropy_mask.extend(other.batch_entropy_mask)
+        self.batch_credits.extend(other.batch_credits)
+        self.batch_engine_log_probs.extend(other.batch_engine_log_probs)
+        self.batch_timesteps.extend(other.batch_timesteps)
+
+
+timestep = int

src_code_for_reproducibility/utils/get_stochastic_game_lengths.py

@@ -0,0 +1,30 @@
+import numpy as np
+
+def get_stochastic_game_lengths(
+    max_length, 
+    nb_games, 
+    continuation_prob, 
+    same_length_batch=False
+):
+    """
+    Generates stochastic game lengths based on a geometric distribution.
+
+    Args:
+        max_length (int): The maximum length a game can have.
+        nb_games (int): The number of games to generate lengths for.
+        continuation_prob (float): The probability of the game continuing after each round.
+        same_length_batch (bool): If True, all games will have the same length.
+
+    Returns:
+        Array: An array of game lengths.
+    """
+    if continuation_prob == 1:
+        return [max_length] * nb_games
+    if same_length_batch:
+        length = np.random.geometric(1 - continuation_prob, 1)
+        game_lengths = np.repeat(length, nb_games)
+    else:
+        game_lengths = np.random.geometric(1 - continuation_prob, nb_games)
+
+    game_lengths = np.where(game_lengths > max_length, max_length, game_lengths)
+    return game_lengths.tolist()

src_code_for_reproducibility/utils/resource_context.py

@@ -0,0 +1,78 @@
+import logging
+import time
+from contextlib import contextmanager
+
+import torch
+
+
+def vram_usage():
+    output = ""
+    for i in range(torch.cuda.device_count()):
+        gpu_memory_allocated = torch.cuda.memory_allocated(i) / (
+            1024**3
+        )  # Convert bytes to GB
+        gpu_memory_reserved = torch.cuda.memory_reserved(i) / (
+            1024**3
+        )  # Convert bytes to GB
+        output += f"GPU {i}: Memory Allocated: {gpu_memory_allocated:.2f} GB, Memory Reserved: {gpu_memory_reserved:.2f} GB"
+    return output
+
+
+def ram_usage():
+    import psutil
+
+    process = psutil.Process()
+    memory_info = process.memory_info()
+    ram_used = memory_info.rss / (1024**3)  # Convert bytes to GB
+    return f"RAM Usage: {ram_used:.2f} GB"
+
+
+@contextmanager
+def resource_logger_context(logger: logging.Logger, task_description: str):
+    """
+    Context manager to log the resource usage of the current task.
+    Args:
+        logger: The logger to use to log the resource usage.
+        task_description: The description of the task to log.
+    Returns:
+        None
+    """
+    try:
+        initial_time = time.time()
+        # Assume CUDA is available and use device 0 only
+        total_mem_bytes = torch.cuda.get_device_properties(0).total_memory
+        initial_total_bytes = (
+            torch.cuda.memory_allocated(0) + torch.cuda.memory_reserved(0)
+        )
+        torch.cuda.reset_peak_memory_stats(0)
+        yield None
+    finally:
+        final_time = time.time()
+        # Ensure kernels within the block are accounted for
+        torch.cuda.synchronize()
+
+        # Compute metrics
+        final_allocated_bytes = torch.cuda.memory_allocated(0)
+        final_reserved_bytes = torch.cuda.memory_reserved(0)
+        final_total_bytes = final_allocated_bytes + final_reserved_bytes
+
+        delta_vram_percent_total = (
+            100 * (final_total_bytes - initial_total_bytes) / total_mem_bytes
+            if total_mem_bytes
+            else 0.0
+        )
+        current_percent_vram_taken = (
+            100 * final_total_bytes / total_mem_bytes if total_mem_bytes else 0.0
+        )
+        block_peak_percent = (
+            100 * torch.cuda.max_memory_allocated(0) / total_mem_bytes
+            if total_mem_bytes
+            else 0.0
+        )
+        delta_time_str = time.strftime(
+            '%H:%M:%S', time.gmtime(final_time - initial_time)
+        )
+
+        logger.info(
+            f"For task: {task_description}, ΔVRAM % (total): {delta_vram_percent_total:.2f}%, Current % of VRAM taken: {current_percent_vram_taken:.2f}%, Block Peak % of device VRAM: {block_peak_percent:.2f}%, ΔTime: {delta_time_str}"
+        )

src_code_for_reproducibility/utils/rollout_tree_chat_htmls.py

@@ -0,0 +1,1921 @@
+from pathlib import Path
+from typing import List
+
+from mllm.utils.rollout_tree_gather_utils import *
+
+
+def html_from_chat_turns(chat_turns: List[ChatTurnLog]) -> str:
+    """
+    Render chat turns as a single, wrapping sequence of messages in time order.
+    Keep badge and message bubble styles, include time on every badge and
+    include rewards on assistant badges. Each message is individually
+    hide/show by click; when hidden, only the badge remains and "(...)" is
+    shown inline (not inside a bubble).
+    """
+    import html
+    import re as _re
+
+    # Prepare ordering: sort by (time_step, original_index) to keep stable order within same step
+    indexed_turns = list(enumerate(chat_turns))
+    indexed_turns.sort(key=lambda t: (t[1].time_step, t[0]))
+    assistant_agents = sorted({t.agent_id for t in chat_turns if t.role == "assistant"})
+    enable_split_view = len(assistant_agents) == 2
+
+    # CSS styles (simplified layout; no time-step or agent-column backgrounds)
+    css = """
+    <style>
+        :root {
+            --font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+            --bg: #ffffff;
+            --text: #1c0b00;
+            --muted-text: #2C3E50;
+            --accent-muted: #BDC3C7;
+            --accent-muted-2: #D0D7DE;
+            --panel-bg: #F8FAFC;
+            --reward-color: #3a2e00; /* dark text for reward pill */
+            --font-size: 14px;
+            --border-width: 2px;
+            --corner-radius: 6px;
+            --pill-radius-left: 999px 0 0 999px;
+            --pill-radius-right: 0 999px 999px 0;
+            --inset-shadow: 0 1px 0 rgba(0,0,0,0.03) inset;
+            
+            /* Chat View Colors */
+            --alice-bg: #dcf8c6;
+            --alice-border: #0eb224;
+            --bob-bg: #ffe4cc;
+            --bob-border: #ef8323;
+            --user-bg: #f5f5f5;
+            --chat-bg: #ffffff;
+        }
+        body {
+            font-family: var(--font-family);
+            margin: 12px;
+            background-color: var(--bg);
+            color: var(--text);
+            font-size: var(--font-size);
+            line-height: 1.5;
+        }
+        
+        /* Chat View Styles */
+        #flow-chat {
+            max-width: 900px;
+            margin: 0 auto;
+            background: var(--chat-bg);
+            padding: 12px 16px 12px 8px;
+            border-radius: 8px;
+        }
+        
+        .simultaneous-messages {
+            display: flex !important;
+            flex-direction: row !important;
+            flex-wrap: nowrap !important;
+            gap: 8px;
+            margin-bottom: 4px;
+            align-items: flex-start;
+            width: 100%;
+            overflow: hidden;
+            box-sizing: border-box;
+        }
+        
+        .simultaneous-messages .chat-message {
+            flex: 1 1 0 !important;
+            margin-bottom: 0 !important;
+            display: flex !important;
+            flex-direction: row !important;
+            align-items: flex-start !important;
+            margin-left: 0 !important;
+            min-width: 0 !important;
+            max-width: 50% !important;
+            gap: 0 !important;
+            overflow: hidden !important;
+        }
+        
+        .simultaneous-messages .chat-message-content {
+            max-width: 100% !important;
+            width: 100%;
+            align-items: flex-start !important;
+            margin-left: 0 !important;
+            overflow: hidden !important;
+        }
+        
+        .simultaneous-messages .chat-message.agent-alice {
+            justify-content: flex-start !important;
+        }
+        
+        .simultaneous-messages .chat-message.agent-bob {
+            justify-content: flex-end !important;
+        }
+        
+        .simultaneous-messages .chat-message.agent-alice .chat-message-content {
+            margin-left: 0 !important;
+            align-items: flex-start !important;
+        }
+        
+        .simultaneous-messages .chat-message.agent-bob .chat-message-content {
+            margin-left: auto !important;
+            margin-right: 0 !important;
+            align-items: flex-end !important;
+        }
+        
+        .simultaneous-messages .chat-bubble {
+            max-width: 100%;
+            word-break: break-word;
+            overflow-wrap: break-word;
+            box-sizing: border-box;
+        }
+        
+        .simultaneous-messages .chat-message.agent-alice .chat-bubble {
+            border-radius: 10px;
+        }
+        
+        .simultaneous-messages .chat-message.agent-bob .chat-bubble {
+            border-radius: 10px;
+        }
+        
+        .simultaneous-messages .chat-message.agent-alice .chat-header {
+            justify-content: flex-start;
+            flex-shrink: 0;
+        }
+        
+        .simultaneous-messages .chat-message.agent-bob .chat-header {
+            justify-content: flex-end;
+            flex-shrink: 0;
+        }
+        
+        .simultaneous-messages .chat-reasoning {
+            max-width: 100%;
+            overflow-wrap: break-word;
+        }
+        
+        .chat-message {
+            display: flex;
+            margin-bottom: 2px;
+            align-items: flex-end;
+            gap: 6px;
+            position: relative;
+            margin-left: 36px;
+        }
+        
+        .chat-message.agent-alice {
+            margin-left: 0;
+        }
+        
+        .chat-message.agent-alice::before {
+            left: 0;
+        }
+        
+        .chat-message.role-user {
+            opacity: 0.7;
+        }
+        
+        .chat-message::before {
+            content: '';
+            position: absolute;
+            left: -36px;
+            top: 0;
+            bottom: 0;
+            width: 36px;
+            pointer-events: auto;
+        }
+        
+        .merge-btn {
+            position: absolute;
+            left: -30px;
+            top: 50%;
+            transform: translateY(-50%);
+            width: 26px;
+            height: 26px;
+            border-radius: 4px;
+            border: 1.5px solid var(--accent-muted);
+            background: white;
+            cursor: pointer;
+            font-size: var(--font-size);
+            opacity: 0;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            transition: opacity 0.2s ease, transform 0.1s ease;
+            padding: 0;
+            line-height: 1;
+            z-index: 10;
+        }
+        
+        .chat-message:hover .merge-btn,
+        .merge-btn:hover {
+            opacity: 1;
+        }
+        
+        .merge-btn:hover {
+            background: var(--panel-bg);
+            border-color: var(--accent-muted-2);
+            transform: translateY(-50%) scale(1.15);
+            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.15);
+        }
+        
+        .merge-btn:active {
+            transform: translateY(-50%) scale(0.95);
+        }
+        
+        .chat-message.agent-alice .merge-btn {
+            left: -30px;
+        }
+        
+        .chat-message.role-user .merge-btn {
+            display: none !important;
+        }
+        
+        .simultaneous-messages .merge-btn {
+            opacity: 0 !important;
+            pointer-events: none;
+        }
+        
+        .simultaneous-messages {
+            padding: 6px 0 6px 0 !important;
+            margin-left: 0 !important;
+            margin-right: 0 !important;
+            position: relative !important;
+            background: transparent !important;
+            border-radius: 0 !important;
+            box-sizing: border-box !important;
+            overflow: visible !important;
+            max-width: 100% !important;
+            border: none !important;
+            transition: padding 0.2s ease !important;
+        }
+        
+        .simultaneous-messages:hover {
+            padding-top: 40px !important;
+        }
+        
+        .simultaneous-messages::before {
+            content: '⇅ Merged';
+            position: absolute;
+            left: 0 !important;
+            top: 8px !important;
+            font-size: var(--font-size);
+            font-weight: 500;
+            color: #888;
+            pointer-events: none;
+            opacity: 0;
+            transition: opacity 0.2s ease;
+        }
+        
+        .simultaneous-messages:hover::before {
+            opacity: 1;
+        }
+        
+        .unmerge-btn {
+            position: absolute !important;
+            right: 0 !important;
+            top: 6px !important;
+            width: 36px !important;
+            height: 28px !important;
+            border-radius: 5px !important;
+            border: 2px solid #d63031 !important;
+            background: white !important;
+            cursor: pointer !important;
+            font-size: var(--font-size) !important;
+            font-weight: bold !important;
+            color: #d63031 !important;
+            display: flex !important;
+            align-items: center !important;
+            justify-content: center !important;
+            transition: all 0.2s ease !important;
+            padding: 0 !important;
+            line-height: 1 !important;
+            z-index: 1000 !important;
+            flex: none !important;
+            pointer-events: auto !important;
+            box-shadow: 0 2px 6px rgba(214, 48, 49, 0.3) !important;
+            opacity: 0 !important;
+        }
+        
+        .simultaneous-messages:hover .unmerge-btn {
+            opacity: 1 !important;
+        }
+        
+        .unmerge-btn:hover {
+            background: #ffe5e5 !important;
+            border-color: #b71c1c !important;
+            transform: scale(1.1) !important;
+            box-shadow: 0 3px 8px rgba(214, 48, 49, 0.4) !important;
+        }
+        
+        .unmerge-btn:active {
+            transform: scale(0.95) !important;
+            background: #ffcccc !important;
+        }
+        
+        .chat-message-content {
+            max-width: 72%;
+            display: flex;
+            flex-direction: column;
+            gap: 2px;
+        }
+        
+        .chat-message.agent-alice .chat-message-content {
+            align-items: flex-start;
+        }
+        
+        .chat-message.agent-bob .chat-message-content {
+            align-items: flex-end;
+            margin-left: auto;
+        }
+        
+        .chat-bubble {
+            padding: 6px 10px;
+            border-radius: 10px;
+            word-wrap: break-word;
+            position: relative;
+            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
+            line-height: 1.4;
+        }
+        
+        .chat-message.agent-alice .chat-bubble {
+            background: var(--alice-bg);
+            border: 2px solid var(--alice-border);
+            border-radius: 10px 10px 10px 2px;
+        }
+        
+        .chat-message.agent-bob .chat-bubble {
+            background: var(--bob-bg);
+            border: 2px solid var(--bob-border);
+            border-radius: 10px 10px 2px 10px;
+        }
+        
+        .chat-message.role-user .chat-bubble {
+            background: var(--user-bg);
+            border: 2px solid #d0d0d0;
+        }
+        
+        .chat-header {
+            display: flex;
+            align-items: center;
+            gap: 4px;
+            margin-bottom: 2px;
+            font-size: var(--font-size);
+            font-weight: 600;
+            line-height: 1.2;
+        }
+        
+        .chat-message.agent-alice .chat-header {
+            color: var(--alice-border);
+        }
+        
+        .chat-message.agent-bob .chat-header {
+            color: var(--bob-border);
+        }
+        
+        .chat-timestamp {
+            font-size: var(--font-size);
+            color: var(--muted-text);
+            margin-top: 1px;
+            opacity: 0.75;
+        }
+        
+        .chat-reward {
+            display: inline-flex;
+            align-items: center;
+            background: linear-gradient(90deg, #fffdf2 0%, #ffffff 75%);
+            color: #000000;
+            font-weight: 600;
+            font-size: var(--font-size);
+            padding: 1px 5px;
+            border-radius: 3px;
+            border: 1px solid #f4e6a8;
+            margin-left: 4px;
+            line-height: 1.3;
+        }
+        
+        .chat-reasoning {
+            font-size: var(--font-size);
+            font-style: italic;
+            color: #555;
+            margin-bottom: 2px;
+            padding: 4px 8px;
+            background: rgba(0, 0, 0, 0.03);
+            border-radius: 5px;
+            cursor: pointer;
+            line-height: 1.3;
+        }
+        
+        .chat-reasoning.collapsed .reasoning-text {
+            display: none;
+        }
+        
+        .chat-reasoning.collapsed::after {
+            content: ' (click to expand)';
+            color: #777;
+        }
+        
+        .chat-group-divider {
+            display: flex;
+            align-items: center;
+            gap: 8px;
+            width: 100%;
+            margin: 8px 0 4px 0;
+            position: relative;
+            cursor: pointer;
+            user-select: none;
+        }
+        
+        .chat-group-divider::before,
+        .chat-group-divider::after {
+            content: "";
+            flex: 1 1 auto;
+            height: 2px;
+            background: linear-gradient(90deg, rgba(224,230,235,0), var(--accent-muted-2) 30%, var(--accent-muted-2) 70%, rgba(224,230,235,0));
+        }
+        
+        .chat-group-label {
+            display: inline-block;
+            background: white;
+            padding: 2px 12px;
+            border-radius: 999px;
+            font-size: var(--font-size);
+            font-weight: 700;
+            color: var(--muted-text);
+            border: 1.5px solid var(--accent-muted);
+            box-shadow: 0 1px 3px rgba(0, 0, 0, 0.08);
+            line-height: 1.4;
+            position: relative;
+            transition: background 0.2s ease;
+        }
+        
+        .chat-group-divider:hover .chat-group-label {
+            background: var(--panel-bg);
+        }
+        
+        .chat-group-label::before {
+            content: '▼ ';
+            font-size: 0.8em;
+            display: inline-block;
+            transition: transform 0.2s ease;
+            opacity: 0;
+        }
+        
+        .chat-group-divider:hover .chat-group-label::before {
+            opacity: 1;
+        }
+        
+        .chat-group-divider.collapsed .chat-group-label::before {
+            content: '▶ ';
+            opacity: 1;
+        }
+        
+        .chat-group-divider.collapsed + * {
+            display: none !important;
+        }
+        
+        /* Hide collapsed rounds in strong hide mode */
+        .strong-hide .chat-group-divider.collapsed {
+            display: none !important;
+        }
+        
+        /* Chat view width control */
+        #flow-chat {
+            --chat-width: 900px;
+            max-width: var(--chat-width);
+            margin: 0 auto;
+        }
+        
+        /* Hide user messages when toggle is on */
+        #flow-chat.hide-user-messages .chat-message.role-user {
+            display: none;
+        }
+        
+        /* Hide rewards when hiding user messages */
+        #flow-chat.hide-user-messages .chat-reward {
+            display: none;
+        }
+        
+        /* Round context annotations */
+        .round-context {
+            text-align: center;
+            margin: 4px auto;
+            max-width: 100%;
+        }
+        
+        .round-context-edit {
+            min-height: 20px;
+            padding: 5px 10px;
+            border: 1.5px dashed var(--accent-muted);
+            border-radius: 6px;
+            background: #fafafa;
+            cursor: text;
+            transition: all 0.2s ease;
+            outline: none;
+            font-size: var(--font-size);
+            line-height: 1.3;
+            user-select: text;
+            -webkit-user-select: text;
+            -moz-user-select: text;
+            -ms-user-select: text;
+        }
+        
+        .round-context-edit:focus {
+            border-style: solid;
+            border-color: var(--accent-muted-2);
+            background: #ffffff;
+            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
+        }
+        
+        .round-context-edit:empty:before {
+            content: attr(data-placeholder);
+            color: #999;
+            font-style: italic;
+        }
+        
+        .round-context-controls {
+            display: none;
+            justify-content: center;
+            gap: 4px;
+            margin-top: 4px;
+            flex-wrap: wrap;
+        }
+        
+        .round-context-edit:focus + .round-context-controls,
+        .round-context-controls:hover,
+        .round-context:focus-within .round-context-controls {
+            display: flex;
+        }
+        
+        .context-color-btn {
+            width: 22px;
+            height: 22px;
+            border-radius: 50%;
+            border: 1.5px solid #fff;
+            box-shadow: 0 1px 2px rgba(0, 0, 0, 0.15);
+            cursor: pointer;
+            transition: transform 0.1s ease;
+        }
+        
+        .context-color-btn:hover {
+            transform: scale(1.15);
+        }
+        
+        .context-color-btn:active {
+            transform: scale(0.95);
+        }
+        
+        /* Split agent context boxes */
+        .split-agent-context {
+            display: flex;
+            gap: 6px;
+            margin: 4px auto;
+            max-width: 100%;
+            align-items: flex-start;
+        }
+        
+        .agent-context-box {
+            flex: 1;
+            min-width: 0;
+            position: relative;
+        }
+        
+        .agent-context-box .round-context-edit {
+            margin: 0;
+            border-radius: 6px;
+            padding: 4px 8px;
+            min-height: 18px;
+        }
+        
+        .agent-context-box.agent-alice .round-context-edit {
+            border-color: var(--alice-border);
+            background: rgba(14, 178, 36, 0.03);
+        }
+        
+        .agent-context-box.agent-bob .round-context-edit {
+            border-color: var(--bob-border);
+            background: rgba(239, 131, 35, 0.03);
+        }
+        
+        .agent-context-box.agent-alice .round-context-edit:focus {
+            border-color: var(--alice-border);
+            box-shadow: 0 2px 8px rgba(14, 178, 36, 0.2);
+            background: rgba(14, 178, 36, 0.05);
+        }
+        
+        .agent-context-box.agent-bob .round-context-edit:focus {
+            border-color: var(--bob-border);
+            box-shadow: 0 2px 8px rgba(239, 131, 35, 0.2);
+            background: rgba(239, 131, 35, 0.05);
+        }
+        
+        .agent-context-box .round-context-edit::before {
+            font-weight: 700;
+            font-size: var(--font-size);
+            margin-right: 5px;
+            letter-spacing: 0.2px;
+        }
+        
+        .agent-context-box.agent-alice .round-context-edit::before {
+            content: 'Alice Prompt Summary:';
+            color: var(--alice-border);
+        }
+        
+        .agent-context-box.agent-bob .round-context-edit::before {
+            content: 'Bob Prompt Summary:';
+            color: var(--bob-border);
+        }
+        
+        /* Empty context boxes will be hidden by JavaScript when strong hide is enabled */
+        .messages-flow { display: block; }
+        .split-wrapper { display: flex; gap: 4px; align-items: flex-start; position: relative; }
+        .split-col { flex:1 1 0; min-width:0; }
+        /* In split view keep same inline density as linear view */
+        .split-col .chat-turn { display: inline; }
+        .split-wrapper.resizing { user-select: none; }
+    .split-resizer { width:4px; cursor: col-resize; flex:0 0 auto; align-self: stretch; position: relative; background: linear-gradient(90deg, rgba(224,230,235,0), var(--accent-muted-2) 30%, var(--accent-muted-2) 70%, rgba(224,230,235,0)); border-radius:2px; transition: background .15s ease, width .15s ease; }
+    .split-resizer:hover { background: linear-gradient(90deg, rgba(224,230,235,0), var(--accent-muted) 35%, var(--accent-muted) 65%, rgba(224,230,235,0)); }
+    .split-resizer.dragging { background: linear-gradient(90deg, rgba(224,230,235,0), var(--accent-muted) 25%, var(--accent-muted) 75%, rgba(224,230,235,0)); }
+    /* Inline reasoning (removed toggle to prevent layout shift on click) */
+    .reasoning-inline { display:inline; font-size:var(--font-size); font-style:italic; color:#555; white-space:pre-wrap; margin-right:4px; cursor:pointer; position:relative; }
+    .reasoning-inline .reasoning-text { display:inline; }
+    .reasoning-inline .reasoning-icon { display:inline-block; margin-right:2px; }
+    .reasoning-inline.collapsed .reasoning-text { display:none; }
+    .reasoning-inline.collapsed::after { content:'(...)'; font-style:italic; color:#777; margin-left:4px; }
+    .message-box .main-content { white-space:normal; }
+        /* tighten spacing */
+        .split-col .group-divider { margin:4px 0 2px 0; }
+        .toolbar {
+            display: flex;
+            align-items: center;
+            gap: 8px;
+            margin-bottom: 0;
+            font-size: var(--font-size);
+            max-height: 0;
+            overflow: hidden;
+            opacity: 0;
+            pointer-events: none;
+            transition: max-height 0.2s ease, opacity 0.2s ease;
+            flex-wrap: wrap;
+        }
+        .toolbar-wrap { position: sticky; top: 0; z-index: 10; background: var(--bg); }
+        .toolbar-hotzone { height: 6px; }
+        .toolbar-wrap:hover .toolbar { max-height: 500px; opacity: 1; pointer-events: auto; margin-bottom: 12px; }
+        .toolbar * { pointer-events: auto !important; }
+        .toolbar input,
+        .toolbar select { z-index: 100 !important; position: relative; }
+        .toolbar input[type="number"],
+        .toolbar input[type="text"],
+        .toolbar select {
+            width: 72px;
+            padding: 2px 6px;
+            border: 1px solid var(--accent-muted);
+            border-radius: var(--corner-radius);
+            background: var(--bg);
+            user-select: text !important;
+            -webkit-user-select: text !important;
+            -moz-user-select: text !important;
+            -ms-user-select: text !important;
+            pointer-events: auto !important;
+            cursor: pointer !important;
+        }
+        .toolbar input[type="text"] {
+            cursor: text !important;
+        }
+        .toolbar input[type="text"]:focus,
+        .toolbar input[type="number"]:focus,
+        .toolbar select:focus {
+            outline: 2px solid #0066cc;
+            outline-offset: 1px;
+        }
+        .toolbar button {
+            padding: 4px 8px;
+            border: 1px solid var(--accent-muted);
+            background: var(--panel-bg);
+            border-radius: var(--corner-radius);
+            cursor: pointer;
+        }
+        .chat-turn {
+            display: inline; /* inline like text */
+            background: transparent;
+            position: relative;
+            cursor: pointer;
+        }
+        /* No agent-specific background distinctions */
+        .turn-content {
+            white-space: normal;
+            color: var(--text);
+            font-size: var(--font-size);
+            display: inline; /* inline flow */
+        }
+        .chat-turn .agent-badge { margin-right: 0; vertical-align: baseline; }
+        .agent-badge {
+            display: inline;
+            position: relative;
+            border: var(--border-width) solid var(--accent-muted); /* slightly thicker */
+            border-radius: var(--pill-radius-left); /* round left and bottom-right */
+            font-size: var(--font-size);
+            color: var(--muted-text);
+            background: var(--panel-bg);
+            box-shadow: var(--inset-shadow);
+            line-height: 1.2;
+            border-right: 0;
+        }
+        /* Use flex on assistant badges to vertically center reward pill */
+        .chat-turn.role-assistant .agent-badge { display: inline-flex; align-items: center; }
+        .agent-badge::after {
+            content: none;
+        }
+        /* removed external separator; emoji is rendered inside message bubble */
+        .agent-name { font-weight: 700; }
+        .emoji-bw { filter: grayscale(100%); opacity: 0.95; font-size: var(--font-size); vertical-align: baseline; margin: 0; position: relative; top: -1px; line-height: 1; display: inline-block; }
+        .ts-badge {
+            position: relative;
+            display: inline;
+            border: var(--border-width) solid var(--accent-muted-2); /* slightly thicker */
+            border-radius: var(--corner-radius); /* not a pill */
+            font-size: var(--font-size);
+            # font-weight: 700;
+            color: var(--muted-text);
+            background: #F4F8FB; /* subtle tint */
+            # padding: 1px 6px; /* slight padding for visibility */
+            margin-right: 8px; /* small gap from following content */
+            pointer-events: auto; /* allow events so we can ignore them in JS */
+        }
+        /* Hide timestep badges when grouping by 1 */
+        .hide-ts-badges .ts-badge { display: none; }
+        /* Strong hide: completely hide collapsed turns */
+        .strong-hide .chat-turn.collapsed { display: none; }
+        .ts-badge::before {
+            content: "";
+            position: relative;
+            background: var(--accent-muted-2);
+            border-radius: 2px;
+        }
+        .agent-badge { margin-left: 6px;  }
+        .message-box {
+            display: inline; /* inline bubble behaving like text */
+            font-size: var(--font-size);
+            border: var(--border-width) solid var(--accent-muted);
+            border-radius: var(--pill-radius-right); /* round left and bottom-right */
+            position: relative;
+            background: var(--bg);
+            vertical-align: baseline;
+            line-height: 1.2;
+            padding-left: 0;
+            border-left: 0;
+        }
+        .chat-turn.agent-alice.role-assistant .message-box::before { color: #0eb224; }
+        .chat-turn.agent-bob.role-assistant .message-box::before { color: #ef8323; }
+        .chat-turn.collapsed .message-box::before { display: none; }
+        /* Assistant bubble border colors by common agent names */
+        .chat-turn.agent-alice.role-assistant .message-box { border-color: #0eb224; }
+        .chat-turn.agent-bob.role-assistant .message-box { border-color: #ef8323; }
+        /* Tie badge and seam to agent color for a cohesive capsule, assistants only */
+    .chat-turn.agent-alice.role-assistant .agent-badge { border-color: #0eb224; background: rgba(14,178,36,0.08); }
+        .chat-turn.agent-alice.role-assistant .agent-badge::after { border-right-color: #0eb224; }
+        .chat-turn.agent-alice.role-assistant .turn-content::before { border-left-color: #0eb224; border-top-color: #0eb224; }
+        .chat-turn.agent-alice.role-assistant .message-box { border-color: #0eb224; }
+
+    .chat-turn.agent-bob.role-assistant .agent-badge { border-color: #ef8323; background: rgba(239,131,35,0.10); }
+        .chat-turn.agent-bob.role-assistant .agent-badge::after { border-right-color: #ef8323; }
+        .chat-turn.agent-bob.role-assistant .turn-content::before { border-left-color: #ef8323; border-top-color: #ef8323; }
+        .chat-turn.agent-bob.role-assistant .message-box { border-color: #ef8323; }
+        /* No colored agent-name; keep neutral */
+        .reward {
+            display: inline-flex;
+            align-items: center;
+            justify-content: center;
+            background: linear-gradient(90deg, #fffdf2 0%, #ffffff 75%);
+            color: #000000; /* full black */
+            font-weight: 600; /* slightly bolder */
+            font-family: "Inter", ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Oxygen, Ubuntu, Cantarell, "Fira Sans", "Droid Sans", "Helvetica Neue", Arial, "Noto Sans", sans-serif;
+            font-size: var(--font-size);
+            letter-spacing: 0.15px;
+            line-height: 1;
+            padding: 0 4px 1px 4px; /* slight bottom pad for optical centering */
+            border-radius: 4px;
+            border: 1px solid #f4e6a8;
+            margin: 0 4px;
+            box-shadow: 0 0 0 1px rgba(255,255,255,0.55) inset, 0 1px 2px rgba(0,0,0,0.04);
+        }
+        .message-placeholder { display: none; color: #7f8c8d; font-style: italic; }
+        .chat-turn.collapsed .message-box { color: transparent; font-size: 0; display: inline-block; }
+        .chat-turn.collapsed .message-box::after { content: "(...)"; color: #7f8c8d; font-style: italic; font-size: var(--font-size); line-height: 1.2; }
+        .chat-turn.collapsed .agent-badge,
+        .chat-turn.collapsed .message-box { opacity: 0.3; }
+        /* Group divider - clearer and pretty */
+        .group-divider {
+            display: flex;
+            align-items: center;
+            gap: 8px;
+            width: 100%;
+            margin: 8px 0 4px 0;
+            position: relative;
+            cursor: pointer;
+            user-select: none;
+        }
+        .group-divider::before,
+        .group-divider::after {
+            content: "";
+            flex: 1 1 auto;
+            height: 2px;
+            background: linear-gradient(90deg, rgba(224,230,235,0), var(--accent-muted-2) 30%, var(--accent-muted-2) 70%, rgba(224,230,235,0));
+        }
+        .group-divider .group-label {
+            display: inline-block;
+            border: 1px solid var(--accent-muted);
+            border-radius: 999px;
+            padding: 2px 10px;
+            font-size: var(--group-label-font-size);
+            font-weight: 700;
+            color: var(--muted-text);
+            background: var(--bg);
+            box-shadow: var(--inset-shadow);
+            position: relative;
+            z-index: 1;
+            transition: background 0.2s ease;
+        }
+        
+        .group-divider:hover .group-label {
+            background: var(--panel-bg);
+        }
+        
+        .group-label::before {
+            content: '▼ ';
+            font-size: 0.8em;
+            display: inline-block;
+            transition: transform 0.2s ease;
+            opacity: 0;
+        }
+        
+        .group-divider:hover .group-label::before {
+            opacity: 1;
+        }
+        
+        .group-divider.collapsed .group-label::before {
+            content: '▶ ';
+            opacity: 1;
+        }
+        
+        /* Hide collapsed rounds in strong hide mode */
+        .strong-hide .group-divider.collapsed {
+            display: none !important;
+        }
+        /* Enhance contrast for print / export */
+        body.split-mode .group-divider::before,
+        body.split-mode .group-divider::after {
+            background: linear-gradient(90deg, rgba(224,230,235,0), var(--accent-muted) 25%, var(--accent-muted) 75%, rgba(224,230,235,0));
+        }
+        .chat-turn .turn-content { position: relative; }
+        .chat-turn .turn-content::before {
+            content: none;
+        }
+        .chat-turn .agent-badge {
+            position: relative;
+        }
+        /* removed absolute-positioned emoji to prevent overlap */
+    </style>
+    """
+
+    # HTML structure
+    html_parts = [
+        "<!DOCTYPE html>",
+        "<html>",
+        "<head>",
+        "<meta charset='UTF-8'>",
+        "<title>Chat Turns</title>",
+        css,
+        "<script>\n"
+        "document.addEventListener('DOMContentLoaded', function() {\n"
+        "  const linearFlow = document.getElementById('flow-linear');\n"
+        "  const splitFlow = document.getElementById('flow-split');\n"
+        "  const chatFlow = document.getElementById('flow-chat');\n"
+        "  let splitViewOn = false;\n"
+        "  let chatViewOn = true;\n"
+        "  function activeFlows() { return [chatViewOn && chatFlow ? chatFlow : null, splitViewOn && splitFlow ? splitFlow : null, linearFlow].filter(Boolean).filter(f => f.style.display !== 'none'); }\n"
+        "  // State for range filtering and strong hide\n"
+        "  let currentRangeStart = null;\n"
+        "  let currentRangeEnd = null;\n"
+        "  let strongHideOn = false;\n"
+        "  document.body.addEventListener('click', function(e){\n"
+        "    if (e.target.closest('input, textarea, select, button, .round-context-edit, .toolbar')) { return; }\n"
+        "    if (e.target.closest('.ts-badge')) { return; }\n"
+        "    const r = e.target.closest('.reasoning-inline'); if (r) { e.stopPropagation(); r.classList.toggle('collapsed'); return; }\n"
+        "    const turn = e.target.closest('.chat-turn');\n"
+        "    if (turn) { e.stopPropagation(); turn.classList.toggle('collapsed'); }\n"
+        "  });\n"
+        "  // Reasoning handled via <details>, no JS required\n"
+        "  function applyRangeFilter() {\n"
+        "    for (const flow of activeFlows()) {\n"
+        "      const turns = Array.from(flow.querySelectorAll('.chat-turn'));\n"
+        "      for (const el of turns) {\n"
+        "        const t = parseInt(el.getAttribute('data-time-step') || '0', 10);\n"
+        "        const afterStart = (currentRangeStart === null) || (t >= currentRangeStart);\n"
+        "        const beforeEnd = (currentRangeEnd === null) || (t <= currentRangeEnd);\n"
+        "        el.style.display = (afterStart && beforeEnd) ? '' : 'none';\n"
+        "      }\n"
+        "      const dividers = Array.from(flow.querySelectorAll('.group-divider'));\n"
+        "      for (const d of dividers) {\n"
+        "        let anyVisible = false;\n"
+        "        let el = d.nextElementSibling;\n"
+        "        while (el && !el.classList.contains('group-divider')) {\n"
+        "          if (el.classList.contains('chat-turn')) {\n"
+        "            const disp = getComputedStyle(el).display;\n"
+        "            if (disp !== 'none') { anyVisible = true; break; }\n"
+        "          } else if (el.classList.contains('split-wrapper')) {\n"
+        "            // Search descendants for any visible chat-turn\n"
+        "            const turns = Array.from(el.querySelectorAll('.chat-turn'));\n"
+        "            for (const tEl of turns) {\n"
+        "              const disp2 = getComputedStyle(tEl).display;\n"
+        "              if (disp2 !== 'none') { anyVisible = true; break; }\n"
+        "            }\n"
+        "            if (anyVisible) break;\n"
+        "          }\n"
+        "          el = el.nextElementSibling;\n"
+        "        }\n"
+        "        d.style.display = anyVisible ? '' : 'none';\n"
+        "      }\n"
+        "    }\n"
+        "  }\n"
+        "  function applyGrouping(n) {\n"
+        "    function groupContainer(container, n) {\n"
+        "      Array.from(container.querySelectorAll(':scope > .group-divider')).forEach(el => el.remove());\n"
+        "      if (!n || n <= 0) { return; }\n"
+        "      const turns = Array.from(container.querySelectorAll(':scope > .chat-turn'));\n"
+        "      if (turns.length === 0) return;\n"
+        "      const items = Array.from(container.children).filter(el => !el.classList.contains('group-divider'));\n"
+        "      const frag = document.createDocumentFragment();\n"
+        "      let lastGroup = -1;\n"
+        "      for (const el of items) {\n"
+        "        if (!el.classList.contains('chat-turn')) { frag.appendChild(el); continue; }\n"
+        "        const t = parseInt(el.getAttribute('data-time-step') || '0', 10);\n"
+        "        const g = Math.floor(t / n);\n"
+        "        if (g !== lastGroup) {\n"
+        "          const div = document.createElement('div');\n"
+        "          div.className = 'group-divider';\n"
+        "          const label = document.createElement('span');\n"
+        "          label.className = 'group-label';\n"
+        "          const roundIndex = g + 1;\n"
+        "          label.textContent = `Round ${roundIndex}`;\n"
+        "          div.appendChild(label);\n"
+        "          frag.appendChild(div);\n"
+        "          lastGroup = g;\n"
+        "        }\n"
+        "        frag.appendChild(el);\n"
+        "      }\n"
+        "      container.innerHTML = '';\n"
+        "      container.appendChild(frag);\n"
+        "      container.classList.toggle('hide-ts-badges', n === 1);\n"
+        "      container.classList.toggle('strong-hide', strongHideOn);\n"
+        "    }\n"
+        "    for (const flow of activeFlows()) {\n"
+        "      if (flow.id === 'flow-split') {\n"
+        "        // Snapshot original turns once to avoid drift on repeated grouping\n"
+        "        const getOriginalTurns = () => {\n"
+        "          if (!flow.dataset.origData) {\n"
+        "            const data = [];\n"
+        "            const cols0 = flow.querySelectorAll('.split-col');\n"
+        "            cols0.forEach(col => {\n"
+        "              const agent = col.getAttribute('data-agent') || '';\n"
+        "              col.querySelectorAll(':scope > .chat-turn').forEach(el => {\n"
+        "                const t = parseInt(el.getAttribute('data-time-step')||'0',10);\n"
+        "                data.push({agent, time:t, html: el.outerHTML});\n"
+        "              });\n"
+        "            });\n"
+        "            flow.dataset.origData = JSON.stringify(data);\n"
+        "          }\n"
+        "          return JSON.parse(flow.dataset.origData);\n"
+        "        };\n"
+        "        const original = getOriginalTurns();\n"
+        "        const agents = Array.from(new Set(original.map(o => o.agent))).sort();\n"
+        "        const groups = new Map();\n"
+        "        original.forEach(o => {\n"
+        "          const g = n && n > 0 ? Math.floor(o.time / n) : 0;\n"
+        "          if (!groups.has(g)) groups.set(g, new Map());\n"
+        "          const gm = groups.get(g);\n"
+        "          if (!gm.has(o.agent)) gm.set(o.agent, []);\n"
+        "          gm.get(o.agent).push(o);\n"
+        "        });\n"
+        "        flow.innerHTML = '';\n"
+        "        const sorted = Array.from(groups.keys()).sort((a,b)=>a-b);\n"
+        "        sorted.forEach(g => {\n"
+        "          const div = document.createElement('div');\n"
+        "          div.className = 'group-divider';\n"
+        "          const label = document.createElement('span');\n"
+        "          label.className = 'group-label';\n"
+        "          label.textContent = `Round ${g+1}`;\n"
+        "          div.appendChild(label);\n"
+        "          flow.appendChild(div);\n"
+        "          const wrapper = document.createElement('div');\n"
+        "          wrapper.className = 'split-wrapper';\n"
+        "          agents.forEach(agent => {\n"
+        "            const colDiv = document.createElement('div');\n"
+        "            colDiv.className = 'split-col';\n"
+        "            colDiv.setAttribute('data-agent', agent);\n"
+        "            (groups.get(g).get(agent) || []).forEach(o => { colDiv.insertAdjacentHTML('beforeend', o.html); });\n"
+        "            wrapper.appendChild(colDiv);\n"
+        "          });\n"
+        "          if (wrapper.children.length === 2) { const res = document.createElement('div'); res.className='split-resizer'; wrapper.insertBefore(res, wrapper.children[1]); }\n"
+        "          flow.appendChild(wrapper);\n"
+        "        });\n"
+        "        flow.classList.toggle('hide-ts-badges', n === 1);\n"
+        "        flow.classList.toggle('strong-hide', strongHideOn);\n"
+        "        document.body.classList.add('split-mode');\n"
+        "      } else {\n"
+        "        groupContainer(flow, n);\n"
+        "      }\n"
+        "    }\n"
+        "    applyRangeFilter();\n"
+        "    initSplitResizers();\n"
+        "  }\n"
+        "  function initSplitResizers() {\n"
+        "    const wrappers = document.querySelectorAll('#flow-split .split-wrapper');\n"
+        "    wrappers.forEach(wrap => {\n"
+        "      const resizer = wrap.querySelector('.split-resizer');\n"
+        "      if (!resizer || resizer.dataset.bound) return; resizer.dataset.bound='1';\n"
+        "      const cols = wrap.querySelectorAll('.split-col'); if (cols.length !== 2) return; const c0=cols[0], c1=cols[1];\n"
+        "      c0.style.flex=c1.style.flex='1 1 0'; c0.style.width=c1.style.width='';\n"
+        "      requestAnimationFrame(()=>{ const w0=c0.scrollWidth,w1=c1.scrollWidth,total=w0+w1||1; let p0=w0/total,p1=w1/total; const minP=0.25,maxP=0.75; if(p0<minP){p0=minP;p1=1-p0;} else if(p0>maxP){p0=maxP;p1=1-p0;} c0.style.flex='0 0 '+(p0*100).toFixed(2)+'%'; c1.style.flex='0 0 '+(p1*100).toFixed(2)+'%'; });\n"
+        "      let dragging=false,startX=0,startP0=0;\n"
+        "      const onDown=e=>{ dragging=true; startX=e.clientX; wrap.classList.add('resizing'); resizer.classList.add('dragging'); const rect=wrap.getBoundingClientRect(); const w=rect.width; const c0Rect=c0.getBoundingClientRect(); startP0=c0Rect.width/w; document.body.style.cursor='col-resize'; e.preventDefault(); };\n"
+        "      const onMove=e=>{ if(!dragging)return; const rect=wrap.getBoundingClientRect(); const w=rect.width; let delta=(e.clientX-startX)/w; let newP0=startP0+delta; const minP=0.15,maxP=0.85; if(newP0<minP)newP0=minP; if(newP0>maxP)newP0=maxP; c0.style.flex='0 0 '+(newP0*100).toFixed(2)+'%'; c1.style.flex='0 0 '+((1-newP0)*100).toFixed(2)+'%'; };\n"
+        "      const onUp=()=>{ if(!dragging)return; dragging=false; wrap.classList.remove('resizing'); resizer.classList.remove('dragging'); document.body.style.cursor=''; };\n"
+        "      resizer.addEventListener('mousedown', onDown); window.addEventListener('mousemove', onMove); window.addEventListener('mouseup', onUp);\n"
+        "      resizer.addEventListener('dblclick', e=>{ if(e.shiftKey){ c0.style.flex=c1.style.flex='1 1 0'; requestAnimationFrame(()=>{ const w0=c0.scrollWidth,w1=c1.scrollWidth,total=w0+w1||1; let p0=w0/total,p1=w1/total; const minP=0.25,maxP=0.75; if(p0<minP){p0=minP;p1=1-p0;} else if(p0>maxP){p0=maxP;p1=1-p0;} c0.style.flex='0 0 '+(p0*100).toFixed(2)+'%'; c1.style.flex='0 0 '+(p1*100).toFixed(2)+'%'; }); } else { c0.style.flex='0 0 50%'; c1.style.flex='0 0 50%'; } });\n"
+        "    });\n"
+        "  }\n"
+        "  initSplitResizers();\n"
+        "  const input = document.getElementById('group-size');\n"
+        "  const btn = document.getElementById('apply-grouping');\n"
+        "  if (btn && input) {\n"
+        "    btn.addEventListener('click', () => { const n = parseInt(input.value || '0', 10); applyGrouping(n); });\n"
+        "    input.addEventListener('keydown', (e) => { if (e.key === 'Enter') { const n = parseInt(input.value || '0', 10); applyGrouping(n); } });\n"
+        "  }\n"
+        "  if (input) { input.value = '1'; applyGrouping(1); }\n"
+        "  const rangeStart = document.getElementById('range-start');\n"
+        "  const rangeEnd = document.getElementById('range-end');\n"
+        "  const rangeBtn = document.getElementById('apply-range');\n"
+        "  if (rangeBtn && rangeStart && rangeEnd) {\n"
+        "    const applyRange = () => {\n"
+        "      const sv = parseInt(rangeStart.value || '', 10);\n"
+        "      const ev = parseInt(rangeEnd.value || '', 10);\n"
+        "      currentRangeStart = Number.isFinite(sv) ? sv : null;\n"
+        "      currentRangeEnd = Number.isFinite(ev) ? ev : null;\n"
+        "      applyRangeFilter();\n"
+        "    };\n"
+        "    rangeBtn.addEventListener('click', applyRange);\n"
+        "    rangeStart.addEventListener('keydown', (e) => { if (e.key === 'Enter') applyRange(); });\n"
+        "    rangeEnd.addEventListener('keydown', (e) => { if (e.key === 'Enter') applyRange(); });\n"
+        "  }\n"
+        "  const strongHideBtn = document.getElementById('toggle-strong-hide');\n"
+        "  const strongHideStateEl = document.getElementById('strong-hide-state');\n"
+        "  if (strongHideBtn) {\n"
+        "    const setLabel = () => { if (strongHideStateEl) { strongHideStateEl.textContent = strongHideOn ? 'On' : 'Off'; } };\n"
+        "    strongHideBtn.addEventListener('click', () => { strongHideOn = !strongHideOn; for (const f of activeFlows()) { f.classList.toggle('strong-hide', strongHideOn); } setLabel(); });\n"
+        "    if (strongHideOn) { for (const f of activeFlows()) { f.classList.add('strong-hide'); } }\n"
+        "    setLabel();\n"
+        "  }\n"
+        "  const splitBtn = document.getElementById('toggle-split-view');\n"
+        "  const splitStateEl = document.getElementById('split-view-state');\n"
+        "  if (splitBtn && splitFlow && linearFlow) {\n"
+        "    const updateSplit = () => { if (splitStateEl) splitStateEl.textContent = splitViewOn ? 'On' : 'Off'; };\n"
+        "    splitBtn.addEventListener('click', () => { if (chatViewOn) return; splitViewOn = !splitViewOn; linearFlow.style.display = splitViewOn ? 'none' : ''; splitFlow.style.display = splitViewOn ? '' : 'none'; applyGrouping(parseInt(input.value||'1',10)); updateSplit(); });\n"
+        "    updateSplit();\n"
+        "  }\n"
+        "  const chatBtn = document.getElementById('toggle-chat-view');\n"
+        "  const chatStateEl = document.getElementById('chat-view-state');\n"
+        "  const hideUserBtn = document.getElementById('toggle-hide-user-messages');\n"
+        "  const hideUserStateEl = document.getElementById('hide-user-state');\n"
+        "  const widthControl = document.getElementById('chat-width-control');\n"
+        "  const widthSlider = document.getElementById('chat-width-slider');\n"
+        "  const widthValue = document.getElementById('chat-width-value');\n"
+        "  let hideUserMessages = false;\n"
+        "  if (chatBtn && chatFlow && linearFlow) {\n"
+        "    const updateChat = () => {\n"
+        "      if (chatStateEl) chatStateEl.textContent = chatViewOn ? 'On' : 'Off';\n"
+        "      if (hideUserBtn) hideUserBtn.style.display = chatViewOn ? '' : 'none';\n"
+        "      if (widthControl) widthControl.style.display = chatViewOn ? '' : 'none';\n"
+        "    };\n"
+        "    chatBtn.addEventListener('click', () => {\n"
+        "      chatViewOn = !chatViewOn;\n"
+        "      if (chatViewOn) {\n"
+        "        splitViewOn = false;\n"
+        "        linearFlow.style.display = 'none';\n"
+        "        if (splitFlow) splitFlow.style.display = 'none';\n"
+        "        chatFlow.style.display = '';\n"
+        "        if (splitStateEl) splitStateEl.textContent = 'Off';\n"
+        "      } else {\n"
+        "        chatFlow.style.display = 'none';\n"
+        "        linearFlow.style.display = '';\n"
+        "      }\n"
+        "      updateChat();\n"
+        "    });\n"
+        "    updateChat();\n"
+        "  }\n"
+        "  if (hideUserBtn && hideUserStateEl && chatFlow) {\n"
+        "    const updateHideUser = () => { hideUserStateEl.textContent = hideUserMessages ? 'On' : 'Off'; };\n"
+        "    hideUserBtn.addEventListener('click', () => {\n"
+        "      hideUserMessages = !hideUserMessages;\n"
+        "      chatFlow.classList.toggle('hide-user-messages', hideUserMessages);\n"
+        "      updateHideUser();\n"
+        "    });\n"
+        "    updateHideUser();\n"
+        "  }\n"
+        "  if (widthSlider && widthValue && chatFlow) {\n"
+        "    const savedWidth = localStorage.getItem('chat-view-width');\n"
+        "    if (savedWidth) {\n"
+        "      widthSlider.value = savedWidth;\n"
+        "      chatFlow.style.setProperty('--chat-width', savedWidth + 'px');\n"
+        "      widthValue.textContent = savedWidth + 'px';\n"
+        "    }\n"
+        "    widthSlider.addEventListener('input', (e) => {\n"
+        "      const width = e.target.value;\n"
+        "      chatFlow.style.setProperty('--chat-width', width + 'px');\n"
+        "      widthValue.textContent = width + 'px';\n"
+        "      localStorage.setItem('chat-view-width', width);\n"
+        "    });\n"
+        "  }\n"
+        "  const fontFamilySelect = document.getElementById('font-family-select');\n"
+        "  const fontSizeInput = document.getElementById('font-size-input');\n"
+        "  if (fontFamilySelect) {\n"
+        "    const savedFont = localStorage.getItem('render-font-family');\n"
+        "    if (savedFont) {\n"
+        "      fontFamilySelect.value = savedFont;\n"
+        "      document.body.style.setProperty('--font-family', savedFont);\n"
+        "    }\n"
+        "    fontFamilySelect.addEventListener('change', (e) => {\n"
+        "      const font = e.target.value;\n"
+        "      document.body.style.setProperty('--font-family', font);\n"
+        "      localStorage.setItem('render-font-family', font);\n"
+        "    });\n"
+        "  }\n"
+        "  if (fontSizeInput) {\n"
+        "    const savedSize = localStorage.getItem('render-font-size');\n"
+        "    if (savedSize) {\n"
+        "      fontSizeInput.value = savedSize;\n"
+        "      document.body.style.setProperty('--font-size', savedSize + 'px');\n"
+        "    }\n"
+        "    fontSizeInput.addEventListener('input', (e) => {\n"
+        "      const size = e.target.value;\n"
+        "      document.body.style.setProperty('--font-size', size + 'px');\n"
+        "      localStorage.setItem('render-font-size', size);\n"
+        "    });\n"
+        "  }\n"
+        "  const aliceEmojiInput = document.getElementById('alice-emoji-input');\n"
+        "  const aliceNameInput = document.getElementById('alice-name-input');\n"
+        "  const bobEmojiInput = document.getElementById('bob-emoji-input');\n"
+        "  const bobNameInput = document.getElementById('bob-name-input');\n"
+        "  const applyAgentNamesBtn = document.getElementById('apply-agent-names');\n"
+        "  function loadAgentNames() {\n"
+        "    if (aliceEmojiInput && aliceNameInput && bobEmojiInput && bobNameInput) {\n"
+        "      const savedAliceEmoji = localStorage.getItem('alice-emoji') || '🤖';\n"
+        "      const savedAliceName = localStorage.getItem('alice-name') || 'Alice';\n"
+        "      const savedBobEmoji = localStorage.getItem('bob-emoji') || '🤖';\n"
+        "      const savedBobName = localStorage.getItem('bob-name') || 'Bob';\n"
+        "      aliceEmojiInput.value = savedAliceEmoji;\n"
+        "      aliceNameInput.value = savedAliceName;\n"
+        "      bobEmojiInput.value = savedBobEmoji;\n"
+        "      bobNameInput.value = savedBobName;\n"
+        "      applyAgentNamesToDOM(savedAliceEmoji, savedAliceName, savedBobEmoji, savedBobName);\n"
+        "    }\n"
+        "  }\n"
+        "  function applyAgentNamesToDOM(aliceEmoji, aliceName, bobEmoji, bobName) {\n"
+        "    const agentMap = { 'alice': { name: aliceName, emoji: aliceEmoji }, 'bob': { name: bobName, emoji: bobEmoji } };\n"
+        "    document.querySelectorAll('[data-agent-id]').forEach(el => {\n"
+        "      const agentId = el.getAttribute('data-agent-id');\n"
+        "      if (!agentMap[agentId]) return;\n"
+        "      if (el.classList.contains('agent-name')) {\n"
+        "        el.textContent = agentMap[agentId].name;\n"
+        "      } else if (el.classList.contains('emoji-bw')) {\n"
+        "        const currentEmoji = el.textContent.trim();\n"
+        "        if (currentEmoji === '🤖' || currentEmoji === '👤') {\n"
+        "          el.textContent = agentMap[agentId].emoji;\n"
+        "        }\n"
+        "      }\n"
+        "    });\n"
+        "    const style = document.createElement('style');\n"
+        "    style.id = 'dynamic-agent-names-style';\n"
+        "    const existingStyle = document.getElementById('dynamic-agent-names-style');\n"
+        "    if (existingStyle) existingStyle.remove();\n"
+        "    style.textContent = `\n"
+        "      .agent-context-box.agent-alice .round-context-edit::before {\n"
+        "        content: '${aliceName} Prompt Summary:';\n"
+        "      }\n"
+        "      .agent-context-box.agent-bob .round-context-edit::before {\n"
+        "        content: '${bobName} Prompt Summary:';\n"
+        "      }\n"
+        "    `;\n"
+        "    document.head.appendChild(style);\n"
+        "  }\n"
+        "  if (applyAgentNamesBtn && aliceEmojiInput && aliceNameInput && bobEmojiInput && bobNameInput) {\n"
+        "    [aliceEmojiInput, aliceNameInput, bobEmojiInput, bobNameInput].forEach(input => {\n"
+        "      input.style.pointerEvents = 'auto';\n"
+        "      if (input.tagName === 'INPUT') {\n"
+        "        input.style.userSelect = 'text';\n"
+        "        input.style.webkitUserSelect = 'text';\n"
+        "        input.readOnly = false;\n"
+        "      }\n"
+        "      input.disabled = false;\n"
+        "      const stopAll = (e) => { e.stopPropagation(); e.stopImmediatePropagation(); };\n"
+        "      input.addEventListener('mousedown', stopAll, true);\n"
+        "      input.addEventListener('mouseup', stopAll, true);\n"
+        "      input.addEventListener('click', stopAll, true);\n"
+        "      input.addEventListener('dblclick', stopAll, true);\n"
+        "      input.addEventListener('focus', stopAll, true);\n"
+        "      input.addEventListener('blur', stopAll, true);\n"
+        "      input.addEventListener('paste', stopAll, true);\n"
+        "      input.addEventListener('cut', stopAll, true);\n"
+        "      input.addEventListener('copy', stopAll, true);\n"
+        "      input.addEventListener('select', stopAll, true);\n"
+        "      input.addEventListener('selectstart', stopAll, true);\n"
+        "      input.addEventListener('keydown', stopAll, true);\n"
+        "      input.addEventListener('keyup', stopAll, true);\n"
+        "      input.addEventListener('keypress', stopAll, true);\n"
+        "      input.addEventListener('input', stopAll, true);\n"
+        "      input.addEventListener('change', stopAll, true);\n"
+        "      input.addEventListener('contextmenu', stopAll, true);\n"
+        "    });\n"
+        "    const applyNames = () => {\n"
+        "      const aliceEmoji = aliceEmojiInput.value || '🤖';\n"
+        "      const aliceName = aliceNameInput.value.trim() || 'Alice';\n"
+        "      const bobEmoji = bobEmojiInput.value || '🤖';\n"
+        "      const bobName = bobNameInput.value.trim() || 'Bob';\n"
+        "      localStorage.setItem('alice-emoji', aliceEmoji);\n"
+        "      localStorage.setItem('alice-name', aliceName);\n"
+        "      localStorage.setItem('bob-emoji', bobEmoji);\n"
+        "      localStorage.setItem('bob-name', bobName);\n"
+        "      applyAgentNamesToDOM(aliceEmoji, aliceName, bobEmoji, bobName);\n"
+        "    };\n"
+        "    applyAgentNamesBtn.addEventListener('click', applyNames);\n"
+        "    [aliceNameInput, bobNameInput].forEach(input => {\n"
+        "      input.addEventListener('keydown', (e) => {\n"
+        "        if (e.key === 'Enter') {\n"
+        "          e.preventDefault();\n"
+        "          e.stopPropagation();\n"
+        "          e.stopImmediatePropagation();\n"
+        "          applyNames();\n"
+        "        }\n"
+        "      }, true);\n"
+        "    });\n"
+        "    [aliceEmojiInput, bobEmojiInput].forEach(select => {\n"
+        "      select.addEventListener('change', applyNames);\n"
+        "    });\n"
+        "  }\n"
+        "  loadAgentNames();\n"
+        "  function setupRoundCollapse() {\n"
+        "    document.addEventListener('click', function(e) {\n"
+        "      if (e.target.closest('input, textarea, select, button, .round-context-edit, .toolbar')) { return; }\n"
+        "      const divider = e.target.closest('.chat-group-divider, .group-divider');\n"
+        "      if (!divider) return;\n"
+        "      divider.classList.toggle('collapsed');\n"
+        "      const isCollapsed = divider.classList.contains('collapsed');\n"
+        "      let nextElement = divider.nextElementSibling;\n"
+        "      while (nextElement) {\n"
+        "        if (nextElement.classList.contains('chat-group-divider') || nextElement.classList.contains('group-divider')) {\n"
+        "          break;\n"
+        "        }\n"
+        "        if (isCollapsed) {\n"
+        "          if (!nextElement.dataset.originalDisplay) {\n"
+        "            nextElement.dataset.originalDisplay = nextElement.style.display || getComputedStyle(nextElement).display;\n"
+        "          }\n"
+        "          nextElement.style.display = 'none';\n"
+        "        } else {\n"
+        "          if (nextElement.dataset.originalDisplay) {\n"
+        "            const originalDisplay = nextElement.dataset.originalDisplay;\n"
+        "            nextElement.style.display = originalDisplay === 'none' ? '' : originalDisplay;\n"
+        "            if (nextElement.style.display === originalDisplay && originalDisplay !== 'none') {\n"
+        "              nextElement.style.display = '';\n"
+        "            }\n"
+        "            delete nextElement.dataset.originalDisplay;\n"
+        "          } else {\n"
+        "            nextElement.style.display = '';\n"
+        "          }\n"
+        "        }\n"
+        "        nextElement = nextElement.nextElementSibling;\n"
+        "      }\n"
+        "      e.stopPropagation();\n"
+        "    });\n"
+        "  }\n"
+        "  setupRoundCollapse();\n"
+        "  const strongHideBtnChat = document.getElementById('toggle-strong-hide');\n"
+        "  function applyStrongHideToChat() {\n"
+        "    if (!chatFlow) return;\n"
+        "    chatFlow.classList.toggle('strong-hide', strongHideOn);\n"
+        "    const contextEdits = chatFlow.querySelectorAll('.round-context-edit');\n"
+        "    contextEdits.forEach(edit => {\n"
+        "      const parent = edit.closest('.round-context, .agent-context-box, .split-agent-context');\n"
+        "      if (parent) {\n"
+        "        if (strongHideOn && edit.textContent.trim() === '') {\n"
+        "          parent.style.display = 'none';\n"
+        "        } else {\n"
+        "          parent.style.display = '';\n"
+        "        }\n"
+        "      }\n"
+        "    });\n"
+        "    const splitContexts = chatFlow.querySelectorAll('.split-agent-context');\n"
+        "    splitContexts.forEach(split => {\n"
+        "      if (strongHideOn) {\n"
+        "        const boxes = split.querySelectorAll('.agent-context-box');\n"
+        "        const allEmpty = Array.from(boxes).every(box => {\n"
+        "          const edit = box.querySelector('.round-context-edit');\n"
+        "          return edit && edit.textContent.trim() === '';\n"
+        "        });\n"
+        "        if (allEmpty) split.style.display = 'none';\n"
+        "      }\n"
+        "    });\n"
+        "  }\n"
+        "  if (strongHideBtnChat && chatFlow) {\n"
+        "    strongHideBtnChat.addEventListener('click', () => {\n"
+        "      setTimeout(() => applyStrongHideToChat(), 0);\n"
+        "    });\n"
+        "  }\n"
+        "  document.addEventListener('click', function(e) {\n"
+        "    if (e.target.closest('input, textarea, select, .round-context-edit, .toolbar')) { return; }\n"
+        "    const chatReasoning = e.target.closest('.chat-reasoning');\n"
+        "    if (chatReasoning) {\n"
+        "      chatReasoning.classList.toggle('collapsed');\n"
+        "    }\n"
+        "  });\n"
+        "  function applyColorToSelection(color, element) {\n"
+        "    const selection = window.getSelection();\n"
+        "    if (!selection.rangeCount) return false;\n"
+        "    const range = selection.getRangeAt(0);\n"
+        "    if (!element.contains(range.commonAncestorContainer)) return false;\n"
+        "    const selectedText = range.toString();\n"
+        "    if (!selectedText) return false;\n"
+        "    if (color === 'default') {\n"
+        "      // Remove styling - just extract the text content\n"
+        "      const textNode = document.createTextNode(selectedText);\n"
+        "      range.deleteContents();\n"
+        "      range.insertNode(textNode);\n"
+        "    } else {\n"
+        "      const span = document.createElement('span');\n"
+        "      span.style.color = color;\n"
+        "      span.style.fontWeight = '600';\n"
+        "      try {\n"
+        "        range.surroundContents(span);\n"
+        "      } catch (e) {\n"
+        "        const contents = range.extractContents();\n"
+        "        span.appendChild(contents);\n"
+        "        range.insertNode(span);\n"
+        "      }\n"
+        "    }\n"
+        "    return true;\n"
+        "  }\n"
+        "  let lastFocusedContextEdit = null;\n"
+        "  document.addEventListener('focusin', function(e) {\n"
+        "    if (e.target.classList.contains('round-context-edit')) {\n"
+        "      lastFocusedContextEdit = e.target;\n"
+        "    }\n"
+        "  });\n"
+        "  document.addEventListener('mousedown', function(e) {\n"
+        "    if (e.target.classList.contains('context-color-btn')) {\n"
+        "      e.preventDefault();\n"
+        "    }\n"
+        "  });\n"
+        "  document.addEventListener('click', function(e) {\n"
+        "    if (e.target.closest('input:not(.round-context-edit), textarea, select') && !e.target.classList.contains('context-color-btn')) { return; }\n"
+        "    if (e.target.classList.contains('context-color-btn')) {\n"
+        "      e.preventDefault();\n"
+        "      const color = e.target.dataset.color;\n"
+        "      const controls = e.target.closest('.round-context-controls');\n"
+        "      const contextEdit = controls ? controls.previousElementSibling : null;\n"
+        "      if (contextEdit && contextEdit.classList.contains('round-context-edit')) {\n"
+        "        contextEdit.focus();\n"
+        "        const selection = window.getSelection();\n"
+        "        if (selection.rangeCount > 0 && selection.toString().length > 0 && contextEdit.contains(selection.anchorNode)) {\n"
+        "          if (applyColorToSelection(color, contextEdit)) {\n"
+        "            const key = contextEdit.dataset.contextKey;\n"
+        "            localStorage.setItem(key, contextEdit.innerHTML);\n"
+        "          }\n"
+        "        } else {\n"
+        "          try {\n"
+        "            if (color !== 'default') {\n"
+        "              document.execCommand('styleWithCSS', false, true);\n"
+        "              document.execCommand('foreColor', false, color);\n"
+        "            }\n"
+        "            const key = contextEdit.dataset.contextKey;\n"
+        "            setTimeout(() => localStorage.setItem(key, contextEdit.innerHTML), 10);\n"
+        "          } catch (e) {\n"
+        "            console.log('Color command failed:', e);\n"
+        "          }\n"
+        "        }\n"
+        "      }\n"
+        "    }\n"
+        "  });\n"
+        "  const contextEdits = document.querySelectorAll('.round-context-edit');\n"
+        "  contextEdits.forEach(edit => {\n"
+        "    edit.addEventListener('input', function() {\n"
+        "      const key = this.dataset.contextKey;\n"
+        "      localStorage.setItem(key, this.innerHTML);\n"
+        "    });\n"
+        "    const key = edit.dataset.contextKey;\n"
+        "    const saved = localStorage.getItem(key);\n"
+        "    if (saved) {\n"
+        "      edit.innerHTML = saved;\n"
+        "    }\n"
+        "  });\n"
+        "  document.addEventListener('click', function(e) {\n"
+        "    if (e.target.closest('input, textarea, select, .round-context-edit') && !e.target.classList.contains('merge-btn') && !e.target.classList.contains('unmerge-btn')) { return; }\n"
+        "    if (e.target.classList.contains('merge-btn')) {\n"
+        "      e.preventDefault();\n"
+        "      e.stopPropagation();\n"
+        "      const msgId = e.target.dataset.msgId;\n"
+        "      const currentMsg = e.target.closest('.chat-message');\n"
+        "      if (!currentMsg) return;\n"
+        "      if (currentMsg.classList.contains('role-user')) {\n"
+        "        alert('Cannot merge user messages');\n"
+        "        return;\n"
+        "      }\n"
+        "      let nextMsg = currentMsg.nextElementSibling;\n"
+        "      while (nextMsg && !nextMsg.classList.contains('chat-message')) {\n"
+        "        nextMsg = nextMsg.nextElementSibling;\n"
+        "      }\n"
+        "      while (nextMsg && nextMsg.classList.contains('role-user')) {\n"
+        "        nextMsg = nextMsg.nextElementSibling;\n"
+        "        while (nextMsg && !nextMsg.classList.contains('chat-message')) {\n"
+        "          nextMsg = nextMsg.nextElementSibling;\n"
+        "        }\n"
+        "      }\n"
+        "      if (!nextMsg || nextMsg.classList.contains('chat-message') === false) {\n"
+        "        alert('No next assistant message to merge with');\n"
+        "        return;\n"
+        "      }\n"
+        "      if (nextMsg.classList.contains('role-user')) {\n"
+        "        alert('Cannot merge with user messages');\n"
+        "        return;\n"
+        "      }\n"
+        "      const parent = currentMsg.parentElement;\n"
+        "      if (parent.classList.contains('simultaneous-messages')) {\n"
+        "        const wrapper = parent;\n"
+        "        currentMsg.style.display = '';\n"
+        "        currentMsg.classList.remove('merged');\n"
+        "        const refNode = wrapper.nextElementSibling;\n"
+        "        parent.parentElement.insertBefore(currentMsg, refNode);\n"
+        "        if (nextMsg.parentElement === wrapper) {\n"
+        "          parent.parentElement.insertBefore(nextMsg, refNode);\n"
+        "        }\n"
+        "        if (wrapper.children.length === 0) {\n"
+        "          wrapper.remove();\n"
+        "        }\n"
+        "      } else {\n"
+        "        const wrapper = document.createElement('div');\n"
+        "        wrapper.className = 'simultaneous-messages';\n"
+        "        const unmergeBtn = document.createElement('button');\n"
+        "        unmergeBtn.className = 'unmerge-btn';\n"
+        "        unmergeBtn.innerHTML = '✕';\n"
+        "        unmergeBtn.title = 'Click to unmerge messages';\n"
+        "        wrapper.appendChild(unmergeBtn);\n"
+        "        wrapper.dataset.firstMsgId = currentMsg.dataset.msgId;\n"
+        "        wrapper.dataset.secondMsgId = nextMsg.dataset.msgId;\n"
+        "        parent.insertBefore(wrapper, currentMsg);\n"
+        "        wrapper.appendChild(currentMsg);\n"
+        "        wrapper.appendChild(nextMsg);\n"
+        "        currentMsg.classList.add('merged');\n"
+        "        nextMsg.classList.add('merged');\n"
+        "      }\n"
+        "    }\n"
+        "    if (e.target.classList.contains('unmerge-btn')) {\n"
+        "      const wrapper = e.target.closest('.simultaneous-messages');\n"
+        "      if (!wrapper) return;\n"
+        "      const parent = wrapper.parentElement;\n"
+        "      const firstMsgId = wrapper.dataset.firstMsgId;\n"
+        "      const secondMsgId = wrapper.dataset.secondMsgId;\n"
+        "      const messages = Array.from(wrapper.querySelectorAll('.chat-message'));\n"
+        "      const refNode = wrapper.nextElementSibling;\n"
+        "      const firstMsg = messages.find(m => m.dataset.msgId === firstMsgId);\n"
+        "      const secondMsg = messages.find(m => m.dataset.msgId === secondMsgId);\n"
+        "      if (firstMsg) {\n"
+        "        firstMsg.classList.remove('merged');\n"
+        "        firstMsg.style.display = '';\n"
+        "        parent.insertBefore(firstMsg, refNode);\n"
+        "      }\n"
+        "      if (secondMsg) {\n"
+        "        secondMsg.classList.remove('merged');\n"
+        "        secondMsg.style.display = '';\n"
+        "        parent.insertBefore(secondMsg, refNode);\n"
+        "      }\n"
+        "      wrapper.remove();\n"
+        "    }\n"
+        "  });\n"
+        "});\n"
+        "</script>",
+        "</head>",
+        "<body>",
+        '<div class="toolbar-wrap">',
+        '<div class="toolbar-hotzone"></div>',
+        '<div class="toolbar">',
+        '<label for="group-size">Group every</label>',
+        '<input id="group-size" type="number" min="0" step="1" value="1" />',
+        "<span>timesteps</span>",
+        '<button id="apply-grouping">Apply</button>',
+        '<span style="margin-left:8px"></span>',
+        '<label for="range-start"><span class="emoji-bw">🔎</span> Range</label>',
+        '<input id="range-start" type="number" step="1" />',
+        "<span>to</span>",
+        '<input id="range-end" type="number" step="1" />',
+        '<button id="apply-range"><span class="emoji-bw">▶︎</span> Apply</button>',
+        '<button id="toggle-strong-hide"><span class="emoji-bw">🗜️</span> Strong Hide: <span id="strong-hide-state">Off</span></button>',
+        (
+            '<button id="toggle-split-view"><span class="emoji-bw">🪟</span> Split View: <span id="split-view-state">Off</span></button>'
+            if enable_split_view
+            else ""
+        ),
+        '<button id="toggle-chat-view"><span class="emoji-bw">💬</span> Chat View: <span id="chat-view-state">On</span></button>',
+        '<button id="toggle-hide-user-messages"><span class="emoji-bw">👁️</span> Hide Prompts: <span id="hide-user-state">Off</span></button>',
+        '<span id="chat-width-control" style="margin-left:8px;">',
+        '<label for="chat-width-slider"><span class="emoji-bw">↔️</span> Width:</label>',
+        '<input id="chat-width-slider" type="range" min="600" max="1600" step="50" value="900" style="width:120px; vertical-align:middle;" />',
+        '<span id="chat-width-value" style="margin-left:4px;">900px</span>',
+        '</span>',
+        '<span style="margin-left:12px;">',
+        '<label for="font-family-select"><span class="emoji-bw">🔤</span> Font:</label>',
+        '<select id="font-family-select" style="padding:2px 6px; border:1px solid var(--accent-muted); border-radius:var(--corner-radius); background:var(--bg);">',
+        '<option value="\'Segoe UI\', Tahoma, Geneva, Verdana, sans-serif">Segoe UI</option>',
+        '<option value="Arial, sans-serif">Arial</option>',
+        '<option value="\'Helvetica Neue\', Helvetica, sans-serif">Helvetica</option>',
+        '<option value="\'Times New Roman\', Times, serif">Times New Roman</option>',
+        '<option value="Georgia, serif">Georgia</option>',
+        '<option value="\'Courier New\', Courier, monospace">Courier New</option>',
+        '<option value="\'Comic Sans MS\', cursive">Comic Sans</option>',
+        '<option value="\'Trebuchet MS\', sans-serif">Trebuchet MS</option>',
+        '<option value="Verdana, sans-serif">Verdana</option>',
+        '<option value="\'Palatino Linotype\', \'Book Antiqua\', Palatino, serif">Palatino</option>',
+        '<option value="\'Lucida Console\', Monaco, monospace">Lucida Console</option>',
+        '</select>',
+        '</span>',
+        '<span style="margin-left:8px;">',
+        '<label for="font-size-input"><span class="emoji-bw">📏</span> Size:</label>',
+        '<input id="font-size-input" type="number" min="8" max="24" step="1" value="14" style="width:50px;" />',
+        '<span>px</span>',
+        '</span>',
+        '<span style="margin-left:12px; display:flex; align-items:center; gap:8px;">',
+        '<label style="font-weight:600;">Agent Names:</label>',
+        '<select id="alice-emoji-input" style="width:65px; padding:2px 6px; border:1px solid var(--accent-muted); border-radius:var(--corner-radius); background:var(--bg);">',
+        '<option value="🤖">🤖 Robot</option>',
+        '<option value="👤">👤 Human</option>',
+        '</select>',
+        '<input id="alice-name-input" type="text" placeholder="Alice" style="width:80px; padding:2px 6px; border:1px solid var(--accent-muted); border-radius:var(--corner-radius); background:var(--bg);" />',
+        '<span style="margin:0 4px;">|</span>',
+        '<select id="bob-emoji-input" style="width:65px; padding:2px 6px; border:1px solid var(--accent-muted); border-radius:var(--corner-radius); background:var(--bg);">',
+        '<option value="🤖">🤖 Robot</option>',
+        '<option value="👤">👤 Human</option>',
+        '</select>',
+        '<input id="bob-name-input" type="text" placeholder="Bob" style="width:80px; padding:2px 6px; border:1px solid var(--accent-muted); border-radius:var(--corner-radius); background:var(--bg);" />',
+        '<button id="apply-agent-names" style="padding:4px 8px; border:1px solid var(--accent-muted); background:var(--panel-bg); border-radius:var(--corner-radius); cursor:pointer;">Apply</button>',
+        '</span>',
+        "</div>",
+        "</div>",
+        '<div id="flow-linear" class="messages-flow" style="display:none">',
+    ]
+
+    last_time_step = None
+    for original_index, turn in indexed_turns:
+        # Build classes
+        agent_class = f"agent-{re.sub('[^a-z0-9_-]', '-', turn.agent_id.lower())}"
+        role_class = f"role-{turn.role}"
+        collapsed_class = " collapsed" if turn.role == "user" else ""
+
+        # Badge content
+        agent_id_clean = html.escape(turn.agent_id).lower()
+        if turn.role == "assistant":
+            name = html.escape(turn.agent_id)
+            emoji = '<span class="emoji-bw" data-agent-id="' + agent_id_clean + '"> 🤖</span>'
+            raw_val = turn.reward
+            if isinstance(raw_val, (int, float)):
+                reward_val = f"{raw_val:.4f}".rstrip("0").rstrip(".")
+                if len(reward_val) > 8:
+                    reward_val = reward_val[:8] + "…"
+            else:
+                reward_val = str(raw_val)
+            # Format: "🤖 Alice • Reward: 5.5556 • 💬 :"
+            badge_inner = (
+                f'{emoji} <span class="agent-name" data-agent-id="{agent_id_clean}">{name}</span>'
+                f' <span class="sep"> • </span><span class="reward">Reward ⚑ = {reward_val}</span>'
+            )
+        else:
+            # For user messages, show "Prompt of {Agent ID}" in the badge
+            name = html.escape(turn.agent_id)
+            # Format (no reward): "Prompt of Alice • "
+            badge_inner = f'Prompt of <span class="agent-name" data-agent-id="{agent_id_clean}">{name}</span> <span class="sep"> • </span>:'
+
+        badge = f'<span class="agent-badge">{badge_inner}</span>'
+
+        # Inline timestep distinction badge at step boundaries (render before first message)
+        ts_badge_html = ""
+        if last_time_step is None or turn.time_step != last_time_step:
+            ts_badge_html = f'<span class="ts-badge">⏱ {turn.time_step}</span>'
+            last_time_step = turn.time_step
+
+        escaped_content = html.escape(turn.content)
+        reasoning_html = ""
+        if turn.reasoning_content:
+            # Normalize reasoning to avoid leading/newline whitespace that creates visual gaps
+            _raw_reasoning = turn.reasoning_content.replace("\r\n", "\n")
+            _raw_reasoning = _re.sub(
+                r"^\s*\n+", "", _raw_reasoning
+            )  # drop leading blank lines
+            _raw_reasoning = _re.sub(
+                r"\*\*(\s*\n\s*)", r"** ", _raw_reasoning
+            )  # newline right after **
+            _raw_reasoning = _re.sub(
+                r"(\s*\n\s*)\*\*", r" **", _raw_reasoning
+            )  # newline right before **
+            escaped_reasoning = html.escape(_raw_reasoning)
+            reasoning_html = f'<span class="reasoning-inline"><span class="reasoning-icon">💭</span><span class="reasoning-text">{escaped_reasoning}</span></span>'
+        collapsed_text = re.sub(r"\s+", " ", escaped_content).strip()
+
+        html_parts.append(
+            f'<div class="chat-turn {agent_class} {role_class}{collapsed_class}" data-time-step="{turn.time_step}">'
+            f'<div class="turn-content {agent_class} {role_class}">{ts_badge_html}{badge}'
+            f'<span class="message-box">{reasoning_html}<span class="main-content">💬 {collapsed_text}</span></span>'
+            f'<span class="message-placeholder">(...)</span>'
+            f"</div>"
+            f"</div>"
+        )
+
+    html_parts.append("</div>")  # close linear flow
+    if enable_split_view:
+        import html as _html_mod
+
+        html_parts.append(
+            '<div id="flow-split" class="messages-flow" style="display:none">'
+        )
+        html_parts.append('<div class="split-wrapper">')
+        # Per-agent columns
+        per_agent_turns = {
+            aid: [t for t in chat_turns if t.agent_id == aid]
+            for aid in assistant_agents
+        }
+        for idx, aid in enumerate(assistant_agents):
+            turns_agent = per_agent_turns[aid]
+            html_parts.append(
+                f'<div class="split-col" data-agent="{_html_mod.escape(aid)}">'
+            )
+            last_ts_agent = None
+            for turn in turns_agent:
+                agent_class = (
+                    f"agent-{re.sub('[^a-z0-9_-]', '-', turn.agent_id.lower())}"
+                )
+                role_class = f"role-{turn.role}"
+                collapsed_class = " collapsed" if turn.role == "user" else ""
+                ts_badge_html = ""
+                if last_ts_agent is None or turn.time_step != last_ts_agent:
+                    ts_badge_html = f'<span class="ts-badge">⏱ {turn.time_step}</span>'
+                    last_ts_agent = turn.time_step
+                esc_content = _html_mod.escape(turn.content)
+                reasoning_html = ""
+                if turn.reasoning_content:
+                    _raw_reasoning = turn.reasoning_content.replace("\r\n", "\n")
+                    _raw_reasoning = _re.sub(r"^\s*\n+", "", _raw_reasoning)
+                    _raw_reasoning = _re.sub(r"\*\*(\s*\n\s*)", r"** ", _raw_reasoning)
+                    _raw_reasoning = _re.sub(r"(\s*\n\s*)\*\*", r" **", _raw_reasoning)
+                    esc_reasoning = _html_mod.escape(_raw_reasoning)
+                    reasoning_html = f'<span class="reasoning-inline"><span class="reasoning-icon">💭</span><span class="reasoning-text">{esc_reasoning}</span></span>'
+                collapsed_text = re.sub(r"\s+", " ", esc_content).strip()
+                agent_id_clean = _html_mod.escape(turn.agent_id).lower()
+                if turn.role == "assistant":
+                    name = _html_mod.escape(turn.agent_id)
+                    emoji = '<span class="emoji-bw" data-agent-id="' + agent_id_clean + '"> 🤖</span>'
+                    raw_val = turn.reward
+                    if isinstance(raw_val, (int, float)):
+                        reward_val = f"{raw_val:.4f}".rstrip("0").rstrip(".")
+                        if len(reward_val) > 8:
+                            reward_val = reward_val[:8] + "…"
+                    else:
+                        reward_val = str(raw_val)
+                    badge_inner = (
+                        f'{emoji} <span class="agent-name" data-agent-id="{agent_id_clean}">{name}</span>'
+                        f' <span class="sep"> • </span><span class="reward">Reward ⚑ : {reward_val}</span>'
+                    )
+                else:
+                    name = _html_mod.escape(turn.agent_id)
+                    badge_inner = f'Prompt of <span class="agent-name" data-agent-id="{agent_id_clean}">{name}</span> <span class="sep"> • </span>:'
+                badge = f'<span class="agent-badge">{badge_inner}</span>'
+                html_parts.append(
+                    f'<div class="chat-turn {agent_class} {role_class}{collapsed_class}" data-time-step="{turn.time_step}">'
+                    f'<div class="turn-content {agent_class} {role_class}">{ts_badge_html}{badge}'
+                    f'<span class="message-box">{reasoning_html}<span class="main-content">💬 {collapsed_text}</span></span>'
+                    f'<span class="message-placeholder">(...)</span>'
+                    f"</div></div>"
+                )
+            html_parts.append("</div>")  # close split col
+        html_parts.append("</div>")  # split-wrapper
+        html_parts.append("</div>")  # flow-split
+    
+    # Add Chat View
+    import html as _html_mod
+    html_parts.append('<div id="flow-chat" class="messages-flow">')
+    
+    # Helper function to add context annotation areas
+    def add_context_area(position: str, time_step: int):
+        context_key = f"round-context-{position}-{time_step}"
+        placeholder = f"Add context {position} round {time_step}..."
+        color_buttons = ""
+        # Add default/reset color button first
+        color_buttons += (
+            f'<div class="context-color-btn" data-color="default" '
+            f'style="background: linear-gradient(135deg, #000 25%, transparent 25%, transparent 75%, #000 75%), '
+            f'linear-gradient(135deg, #000 25%, transparent 25%, transparent 75%, #000 75%); '
+            f'background-size: 4px 4px; background-position: 0 0, 2px 2px; '
+            f'background-color: #fff;" title="Default color"></div>'
+        )
+        for color_name, color_value in [
+            ('red', '#d32f2f'),
+            ('orange', '#f57c00'),
+            ('yellow', '#f9a825'),
+            ('green', '#388e3c'),
+            ('blue', '#1976d2'),
+            ('purple', '#7b1fa2'),
+            ('gray', '#666666'),
+        ]:
+            color_buttons += (
+                f'<div class="context-color-btn" data-color="{color_value}" '
+                f'style="background-color: {color_value};" title="{color_name}"></div>'
+            )
+        
+        html_parts.append(
+            f'<div class="round-context">'
+            f'<div class="round-context-edit" contenteditable="true" spellcheck="true" '
+            f'data-context-key="{context_key}" '
+            f'data-placeholder="{placeholder}"></div>'
+            f'<div class="round-context-controls">{color_buttons}</div>'
+            f'</div>'
+        )
+    
+    # Helper function to add split agent context boxes
+    def add_split_agent_contexts(position: str, time_step: int):
+        color_buttons = ""
+        # Add default/reset color button first
+        color_buttons += (
+            f'<div class="context-color-btn" data-color="default" '
+            f'style="background: linear-gradient(135deg, #000 25%, transparent 25%, transparent 75%, #000 75%), '
+            f'linear-gradient(135deg, #000 25%, transparent 25%, transparent 75%, #000 75%); '
+            f'background-size: 4px 4px; background-position: 0 0, 2px 2px; '
+            f'background-color: #fff;" title="Default color"></div>'
+        )
+        for color_name, color_value in [
+            ('red', '#d32f2f'),
+            ('orange', '#f57c00'),
+            ('yellow', '#f9a825'),
+            ('green', '#388e3c'),
+            ('blue', '#1976d2'),
+            ('purple', '#7b1fa2'),
+            ('gray', '#666666'),
+        ]:
+            color_buttons += (
+                f'<div class="context-color-btn" data-color="{color_value}" '
+                f'style="background-color: {color_value};" title="{color_name}"></div>'
+            )
+        
+        html_parts.append('<div class="split-agent-context">')
+        
+        # Alice box
+        alice_key = f"agent-context-alice-{position}-{time_step}"
+        alice_placeholder = f"..."
+        html_parts.append(
+            f'<div class="agent-context-box agent-alice">'
+            f'<div class="round-context-edit" contenteditable="true" spellcheck="true" '
+            f'data-context-key="{alice_key}" '
+            f'data-placeholder="{alice_placeholder}"></div>'
+            f'<div class="round-context-controls">{color_buttons}</div>'
+            f'</div>'
+        )
+        
+        # Bob box
+        bob_key = f"agent-context-bob-{position}-{time_step}"
+        bob_placeholder = f"..."
+        html_parts.append(
+            f'<div class="agent-context-box agent-bob">'
+            f'<div class="round-context-edit" contenteditable="true" spellcheck="true" '
+            f'data-context-key="{bob_key}" '
+            f'data-placeholder="{bob_placeholder}"></div>'
+            f'<div class="round-context-controls">{color_buttons}</div>'
+            f'</div>'
+        )
+        
+        html_parts.append('</div>')  # split-agent-context
+    
+    last_time_step_chat = None
+    for original_index, turn in indexed_turns:
+        agent_class = f"agent-{re.sub('[^a-z0-9_-]', '-', turn.agent_id.lower())}"
+        role_class = f"role-{turn.role}"
+        
+        # Add time step divider and beginning context
+        if last_time_step_chat is None or turn.time_step != last_time_step_chat:
+            # Add end contexts for previous round (only regular context, not prompt summary)
+            if last_time_step_chat is not None:
+                add_context_area("end", last_time_step_chat)
+            
+            html_parts.append(
+                f'<div class="chat-group-divider">'
+                f'<span class="chat-group-label">⏱ Round {turn.time_step + 1}</span>'
+                f'</div>'
+            )
+            
+            # Add beginning contexts for new round (both context and prompt summary)
+            add_context_area("beginning", turn.time_step)
+            add_split_agent_contexts("beginning", turn.time_step)
+            
+            last_time_step_chat = turn.time_step
+        
+        # Build chat message with merge controls
+        html_parts.append(f'<div class="chat-message {agent_class} {role_class}" data-msg-id="{original_index}">')
+        
+        # Add merge control button
+        html_parts.append(
+            f'<button class="merge-btn" title="Merge with next message" data-msg-id="{original_index}">⇄</button>'
+        )
+        
+        html_parts.append('<div class="chat-message-content">')
+        
+        # Header with agent name and reward (always show reward)
+        agent_id_clean = _html_mod.escape(turn.agent_id).lower()
+        if turn.role == "assistant":
+            name = _html_mod.escape(turn.agent_id)
+            raw_val = turn.reward
+            if isinstance(raw_val, (int, float)):
+                reward_val = f"{raw_val:.4f}".rstrip("0").rstrip(".")
+                if len(reward_val) > 8:
+                    reward_val = reward_val[:8] + "…"
+            else:
+                reward_val = str(raw_val)
+            header_html = (
+                f'<div class="chat-header">'
+                f'<span class="emoji-bw" data-agent-id="{agent_id_clean}">🤖</span> <span class="agent-name" data-agent-id="{agent_id_clean}">{name}</span>'
+                f'<span class="chat-reward">⚑ {reward_val}</span>'
+                f'</div>'
+            )
+        else:
+            name = _html_mod.escape(turn.agent_id)
+            header_html = f'<div class="chat-header">Prompt of <span class="agent-name" data-agent-id="{agent_id_clean}">{name}</span></div>'
+        
+        html_parts.append(header_html)
+        
+        # Reasoning content if present
+        if turn.reasoning_content:
+            _raw_reasoning = turn.reasoning_content.replace("\r\n", "\n")
+            _raw_reasoning = _re.sub(r"^\s*\n+", "", _raw_reasoning)
+            esc_reasoning = _html_mod.escape(_raw_reasoning)
+            html_parts.append(
+                f'<div class="chat-reasoning collapsed">'
+                f'<span class="reasoning-icon">💭</span> '
+                f'<span class="reasoning-text">{esc_reasoning}</span>'
+                f'</div>'
+            )
+        
+        # Message bubble
+        esc_content = _html_mod.escape(turn.content)
+        html_parts.append(f'<div class="chat-bubble">{esc_content}</div>')
+        
+        html_parts.append('</div>')  # chat-message-content
+        html_parts.append('</div>')  # chat-message
+    
+    # Add end contexts for the last round (only regular context, not prompt summary)
+    if last_time_step_chat is not None:
+        add_context_area("end", last_time_step_chat)
+    
+    html_parts.append("</div>")  # flow-chat
+    html_parts.extend(["</body>", "</html>"])
+
+    return "\n".join(html_parts)
+
+
+def export_html_from_rollout_tree(path: Path, outdir: Path, main_only: bool = False):
+    """Process a rollout tree file and generate HTML files for each path.
+    Creates separate HTML files for the main path and each branch path.
+    The main path is saved in the root output directory, while branch paths
+    are saved in a 'branches' subdirectory.
+
+    Args:
+        path: Path to the rollout tree JSON file
+        outdir: Output directory for HTML files
+        main_only: If True, only export the main trajectory (default: False)
+    """
+    root = load_rollout_tree(path)
+    mgid = root.id
+
+    main_path, branch_paths = get_rollout_tree_paths(root)
+
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    # Create branches subdirectory if we have branch paths
+    if not main_only and branch_paths:
+        branches_dir = outdir / f"mgid:{mgid}_branches_html_renders"
+        branches_dir.mkdir(parents=True, exist_ok=True)
+
+    # Generate HTML for the main path
+    chat_turns = gather_all_chat_turns_for_path(main_path)
+    html_content = html_from_chat_turns(chat_turns)
+    output_file = outdir / f"mgid:{mgid}_main_html_render.render.html"
+    with open(output_file, "w", encoding="utf-8") as f:
+        f.write(html_content)
+
+    # Generate HTML for each branch path
+    for path_obj in branch_paths:
+        chat_turns = gather_all_chat_turns_for_path(path_obj)
+
+        html_content = html_from_chat_turns(chat_turns)
+
+        path_id: str = path_obj.id
+        output_filename = f"{path_id}_html_render.render.html"
+
+        output_file = branches_dir / output_filename
+
+        with open(output_file, "w", encoding="utf-8") as f:
+            f.write(html_content)

src_code_for_reproducibility/utils/rollout_tree_stats.py

@@ -0,0 +1,50 @@
+from typing import Any, Callable, List, Tuple
+
+from mllm.markov_games.rollout_tree import RolloutTreeRootNode
+from mllm.markov_games.simulation import SimulationStepLog
+from mllm.utils.rollout_tree_gather_utils import (
+    gather_simulation_step_logs,
+    get_rollout_tree_paths,
+)
+from mllm.utils.stat_pack import StatPack
+
+
+def get_rollout_tree_stat_tally(
+    rollout_tree: RolloutTreeRootNode,
+    metrics: List[Callable[[SimulationStepLog], List[Tuple[str, float]]]],
+) -> StatPack:
+    stat_tally = StatPack()
+    # get simulation step logs
+    node_list = get_rollout_tree_paths(rollout_tree)[0]
+    simulation_step_logs = gather_simulation_step_logs(node_list)
+    for simulation_step_log in simulation_step_logs:
+        for metric in metrics:
+            metric_result = metric(simulation_step_log)
+            if metric_result is not None:
+                for key, value in metric_result:
+                    stat_tally.add_stat(key, value)
+    return stat_tally
+
+
+def get_rollout_tree_mean_stats(
+    rollout_tree: RolloutTreeRootNode, metrics: List[Callable[[SimulationStepLog], Any]]
+) -> StatPack:
+    """Get the mean stats for a rollout tree."""
+    stat_tally = get_rollout_tree_stat_tally(rollout_tree, metrics)
+    return stat_tally.mean()
+
+
+def get_mean_rollout_tree_stats(
+    rollout_trees: List[RolloutTreeRootNode],
+    metrics: List[Callable[[SimulationStepLog], Any]],
+) -> StatPack:
+    """Get the mean stats for a list of rollout trees."""
+    # TODO complete this
+    stat_tallies = [
+        get_rollout_tree_mean_stats(rollout_tree, metrics)
+        for rollout_tree in rollout_trees
+    ]
+    mean_stat_tally = StatPack()
+    for stat_tally in stat_tallies:
+        mean_stat_tally.add_stats(stat_tally)
+    return mean_stat_tally.mean()

src_code_for_reproducibility/utils/short_id_gen.py

@@ -0,0 +1,11 @@
+import uuid
+
+
+def generate_short_id() -> int:
+    """
+    Generates a short unique ID for tracking adapter versions.
+
+    Returns:
+        int: An 8-digit integer ID.
+    """
+    return int(str(uuid.uuid4().int)[:8])

src_code_for_reproducibility/utils/stat_pack.py

@@ -0,0 +1,113 @@
+import csv
+import json
+import os
+import pickle
+from collections import Counter
+from copy import deepcopy
+from locale import strcoll
+from statistics import mean
+from typing import Any, Dict, Iterator, List, Optional, Tuple, TypedDict
+
+import matplotlib.pyplot as plt
+import numpy as np
+from torch.utils.tensorboard import SummaryWriter
+
+plt.style.use(
+    "https://raw.githubusercontent.com/dereckpiche/DedeStyle/refs/heads/main/dedestyle.mplstyle"
+)
+
+import wandb
+
+from . import wandb_utils
+
+
+class StatPack:
+    def __init__(self):
+        self.data = {}
+
+    def add_stat(self, key: str, value: float | int | None):
+        assert (
+            isinstance(value, float) or isinstance(value, int) or value is None
+        ), f"Value {value} is not a valid type"
+        if key not in self.data:
+            self.data[key] = []
+        self.data[key].append(value)
+
+    def add_stats(self, other: "StatPack"):
+        for key in other.keys():
+            self.add_stat(key, other[key])
+
+    def __getitem__(self, key: str):
+        return self.data[key]
+
+    def __setitem__(self, key: str, value: Any):
+        self.data[key] = value
+
+    def __contains__(self, key: str):
+        return key in self.data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __iter__(self):
+        return iter(self.data)
+
+    def keys(self):
+        return self.data.keys()
+
+    def values(self):
+        return self.data.values()
+
+    def items(self):
+        return self.data.items()
+
+    def mean(self):
+        mean_st = StatPack()
+        for key in self.keys():
+            if isinstance(self[key], list):
+                # TODO: exclude None values
+                non_none_values = [v for v in self[key] if v is not None]
+                if non_none_values:
+                    mean_st[key] = np.mean(np.array(non_none_values))
+                else:
+                    mean_st[key] = None
+        return mean_st
+
+    def store_plots(self, folder: str):
+        os.makedirs(folder, exist_ok=True)
+        for key in self.keys():
+            plt.figure(figsize=(10, 5))
+            plt.plot(self[key])
+            plt.title(key)
+            plt.savefig(os.path.join(folder, f"{key}.pdf"))
+            plt.close()
+
+    def store_numpy(self, folder: str):
+        os.makedirs(folder, exist_ok=True)
+        for key in self.keys():
+            # Sanitize filename components (avoid slashes, spaces, etc.)
+            safe_key = str(key).replace(os.sep, "_").replace("/", "_").replace(" ", "_")
+            values = self[key]
+            # Convert None to NaN for numpy compatibility
+            arr = np.array(
+                [(np.nan if (v is None) else v) for v in values], dtype=float
+            )
+            np.save(os.path.join(folder, f"{safe_key}.npy"), arr)
+
+    def store_json(self, folder: str, filename: str = "stats.json"):
+        os.makedirs(folder, exist_ok=True)
+        with open(os.path.join(folder, filename), "w") as f:
+            json.dump(self.data, f, indent=4)
+
+    def store_csv(self, folder: str):
+        os.makedirs(folder, exist_ok=True)
+        for key in self.keys():
+            with open(os.path.join(folder, f"stats.csv"), "w") as f:
+                writer = csv.writer(f)
+                writer.writerow([key] + self[key])
+
+    def store_pickle(self, folder: str):
+        os.makedirs(folder, exist_ok=True)
+        for key in self.keys():
+            with open(os.path.join(folder, f"stats.pkl"), "wb") as f:
+                pickle.dump(self[key], f)

src_code_for_reproducibility/utils/update_start_epoch.py

@@ -0,0 +1,9 @@
+import os
+
+# During run, set hydra.run.dir=./outputs/{folder}
+def update_start_epoch(cfg, output_directory):
+    if cfg["experiment"]["resume_experiment"]:
+        folders = [f for f in os.listdir(output_directory) if f.startswith("iteration_")]
+        iterations = [int(f.split("_")[1]) for f in folders] if folders else [0]
+        cfg["experiment"]["start_epoch"] = max(iterations)
+    return None