Add files using upload-large-folder tool

.hydra/config.yaml

@@ -0,0 +1,240 @@
+experiment:
+  wandb_enabled: true
+  nb_epochs: 1000
+  nb_matches_per_iteration: 128
+  reinit_matches_each_it: true
+  checkpoint_every_n_iterations: 50
+  start_epoch: 0
+  resume_experiment: true
+  base_seed: 0
+  seed_group_size: 1
+  train: true
+  stat_methods_for_live_wandb: mllm.markov_games.ipd.ipd_statistics
+  name: ipd_vanilla_ad_align_no_agent_buffer
+  agent_buffer: false
+  keep_agent_buffer_count: ${lora_count}
+  agent_buffer_recent_k: -1
+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: {}
+  simulation_class_name: IPD
+  simulation_init_args:
+    rounds_per_game: 10
+    cooperate_actions:
+    - <A>
+    defect_actions:
+    - <B>
+    reward: 3
+    punishment: 1
+    temptation: 5
+    sucker: 0
+  agents:
+    0:
+      agent_id: ${agent_0_id}
+      agent_name: Alice
+      agent_class_name: IPDAgent
+      policy_id: base_llm/agent_adapter
+      init_kwargs:
+        max_errors: 0
+        allow_reasoning: false
+        max_reasoning_chars: 0
+        cooperate_string: <A>
+        defect_string: <B>
+        strategy_prompt: ''
+        intro_prompt: "Welcome to an iterated game. You are playing against another\
+          \ agent.\n\nGame Rules:\n1. In each round, both agents simultaneously choose\
+          \ to either play A or B.\n2. The payoffs for each round are as follows:\n\
+          \   - If both agents play A, both receive 3 points.\n   - If both agents\
+          \ play B, both receive 1 points.\n   - If one agent plays A and the other\
+          \ plays B, the player who plays A receives 0 points and the player who plays\
+          \ B receives 5 points.\n3. The goal is to maximize your total points over\
+          \ all rounds.\n4. You will see the previous play of the other player after\
+          \ the completion of a round.\n\nIn order to play A, you must write \"<A>\"\
+          \ and nothing else.\nIn order to play B, you must write \"<B>\" and nothing\
+          \ else.\n\nDuring the entire game, you must only write \"<A>\" or \"<B>\"\
+          . If you write anything else, your response will be considered as an error.\n\
+          \nThe game starts now:\n"
+        goal_prompt: 'Your goal is to maximize your utility.
+
+          '
+    1:
+      agent_id: ${agent_1_id}
+      agent_name: Bob
+      agent_class_name: IPDAgent
+      policy_id: base_llm/agent_adapter
+      init_kwargs:
+        max_errors: 0
+        allow_reasoning: false
+        max_reasoning_chars: 0
+        cooperate_string: <A>
+        defect_string: <B>
+        strategy_prompt: ''
+        intro_prompt: "Welcome to an iterated game. You are playing against another\
+          \ agent.\n\nGame Rules:\n1. In each round, both agents simultaneously choose\
+          \ to either play A or B.\n2. The payoffs for each round are as follows:\n\
+          \   - If both agents play A, both receive 3 points.\n   - If both agents\
+          \ play B, both receive 1 points.\n   - If one agent plays A and the other\
+          \ plays B, the player who plays A receives 0 points and the player who plays\
+          \ B receives 5 points.\n3. The goal is to maximize your total points over\
+          \ all rounds.\n4. You will see the previous play of the other player after\
+          \ the completion of a round.\n\nIn order to play A, you must write \"<A>\"\
+          \ and nothing else.\nIn order to play B, you must write \"<B>\" and nothing\
+          \ else.\n\nDuring the entire game, you must only write \"<A>\" or \"<B>\"\
+          . If you write anything else, your response will be considered as an error.\n\
+          \nThe game starts now:\n"
+        goal_prompt: 'Your goal is to maximize your utility.
+
+          '
+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: false
+      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.01
+      entropy_topk: null
+      entropy_mask_regex: null
+      kl_coeff: 0.0
+      gradient_clipping: 1.0
+      restrict_tokens: null
+      mini_batch_size: 4
+      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.9
+      use_rloo: true
+      enable_tokenwise_logging: false
+      pg_loss_normalization: nb_tokens
+      truncated_importance_sampling_ratio_cap: 2.0
+      reward_normalizing_constant: 5.0
+      ad_align_force_coop_first_step: false
+      ad_align_clipping: null
+      ad_align_gamma: 0.9
+      ad_align_exclude_k_equals_t: true
+      ad_align_use_sign: false
+      ad_align_beta: 0.5
+      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:
+  max_errors: 0
+  allow_reasoning: false
+  max_reasoning_chars: 0
+  cooperate_string: <A>
+  defect_string: <B>
+  strategy_prompt: ''
+  intro_prompt: "Welcome to an iterated game. You are playing against another agent.\n\
+    \nGame Rules:\n1. In each round, both agents simultaneously choose to either play\
+    \ A or B.\n2. The payoffs for each round are as follows:\n   - If both agents\
+    \ play A, both receive 3 points.\n   - If both agents play B, both receive 1 points.\n\
+    \   - If one agent plays A and the other plays B, the player who plays A receives\
+    \ 0 points and the player who plays B receives 5 points.\n3. The goal is to maximize\
+    \ your total points over all rounds.\n4. You will see the previous play of the\
+    \ other player after the completion of a round.\n\nIn order to play A, you must\
+    \ write \"<A>\" and nothing else.\nIn order to play B, you must write \"<B>\"\
+    \ and nothing else.\n\nDuring the entire game, you must only write \"<A>\" or\
+    \ \"<B>\". If you write anything else, your response will be considered as an\
+    \ error.\n\nThe game starts now:\n"
+  goal_prompt: 'Your goal is to maximize your utility.
+
+    '
+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: ipd_vanilla_ad_align_no_agent_buffer.yaml
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.1'
+    cwd: /home/mila/m/mohammed.muqeeth/AdAlignLLM
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /home/mila/m/mohammed.muqeeth/AdAlignLLM/configs
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /network/scratch/m/mohammed.muqeeth/llm_negotiation/2026_03/ipd_vanilla_ad_align_no_agent_buffer
+    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",
+    "k_proj",
+    "q_proj",
+    "down_proj",
+    "v_proj",
+    "o_proj",
+    "gate_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",
+    "k_proj",
+    "q_proj",
+    "down_proj",
+    "v_proj",
+    "o_proj",
+    "gate_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/__init__.py

@@ -0,0 +1,4 @@
+"""
+File: mllm/__init__.py
+Summary: Initializes the multi-agent large language model package namespace.
+"""

src_code_for_reproducibility/markov_games/ipd/Ipd_hard_coded_agents.py

@@ -0,0 +1,76 @@
+"""
+File: mllm/markov_games/ipd/Ipd_hard_coded_agents.py
+Summary: Contains hand-crafted IPD policies used as deterministic baselines.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Tuple
+
+from mllm.markov_games.ipd.ipd_agent import IPDAgent
+from mllm.markov_games.rollout_tree import AgentActLog, ChatTurn
+
+
+@dataclass
+class AlwaysCooperateIPDAgent(IPDAgent):
+    async def act(self, observation) -> Tuple[Any, AgentActLog]:
+        """
+        Always plays the cooperate action, ignoring observation.
+        Returns the configured cooperate_string so the simulation parses it as "C".
+        """
+
+        action = self.cooperate_string
+
+        # Log a minimal, structured chat turn for consistency with other agents
+        turn_text = f"Playing cooperate: {action}"
+        self.state.chat_history.append(
+            ChatTurn(
+                agent_id=self.agent_id,
+                role="assistant",
+                content=turn_text,
+                is_state_end=True,
+            )
+        )
+
+        act_log = AgentActLog(
+            chat_turns=[self.state.chat_history[-1]],
+            info=None,
+        )
+
+        # Advance internal counters similar to IPDAgent semantics
+        self.state.chat_counter = len(self.state.chat_history)
+        self.state.round_nb = observation.round_nb
+
+        return action, act_log
+
+
+@dataclass
+class AlwaysDefectIPDAgent(IPDAgent):
+    async def act(self, observation) -> Tuple[Any, AgentActLog]:
+        """
+        Always plays the defect action, ignoring observation.
+        Returns the configured defect_string so the simulation parses it as "D".
+        """
+
+        action = self.defect_string
+
+        # Log a minimal, structured chat turn for consistency with other agents
+        turn_text = f"Playing defect: {action}"
+        self.state.chat_history.append(
+            ChatTurn(
+                agent_id=self.agent_id,
+                role="assistant",
+                content=turn_text,
+                is_state_end=True,
+            )
+        )
+
+        act_log = AgentActLog(
+            chat_turns=[self.state.chat_history[-1]],
+            info=None,
+        )
+
+        # Advance internal counters similar to IPDAgent semantics
+        self.state.chat_counter = len(self.state.chat_history)
+        self.state.round_nb = observation.round_nb
+
+        return action, act_log

src_code_for_reproducibility/markov_games/ipd/__init__.py

@@ -0,0 +1,11 @@
+"""
+File: mllm/markov_games/ipd/__init__.py
+Summary: Marks the Iterated Prisoner's Dilemma subpackage.
+"""
+
+from .Ipd_hard_coded_agents import AlwaysCooperateIPDAgent, AlwaysDefectIPDAgent
+
+__all__ = [
+    "AlwaysCooperateIPDAgent",
+    "AlwaysDefectIPDAgent",
+]

src_code_for_reproducibility/models/adapter_training_wrapper.py

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

src_code_for_reproducibility/models/inference_backend.py

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

src_code_for_reproducibility/models/inference_backend_dummy.py

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

src_code_for_reproducibility/models/large_language_model_api.py

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

src_code_for_reproducibility/models/large_language_model_local.py

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

src_code_for_reproducibility/models/scalar_critic.py

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

src_code_for_reproducibility/training/tally_rollout.py

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

src_code_for_reproducibility/training/tally_tokenwise.py

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

src_code_for_reproducibility/training/trainer_ad_align.py

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

src_code_for_reproducibility/training/trainer_common.py

@@ -0,0 +1,1032 @@
+"""
+File: mllm/training/trainer_common.py
+Summary: Shared trainer utilities, base classes, and gradient helpers.
+"""
+
+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):
+    """
+    Shared scaffolding for policy-gradient trainers (optimizer wiring, logging, etc.).
+
+    Subclasses implement `set_agent_trajectory_data` / `share_advantage_data`
+    to plug in algorithm-specific behavior.
+    """
+
+    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",
+        no_rloo_grouping: bool = False,
+    ):
+        """
+        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
+        self.no_rloo_grouping = no_rloo_grouping
+
+    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.
+        """
+        # Gradients flow only through the kept logits; masking is recomputed per batch for clarity.
+
+        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 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. Causal LM attention masks are implicit;
+                # wire up a custom mask here only if the policy deviates from standard autoregressive behavior.
+                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 accelerator state so we do not accumulate references between optimizer steps.
+            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:
+        """
+        Compute (and optionally whiten) advantages while training the critic in mini-batches.
+        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 and not self.no_rloo_grouping:
+                    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:
+        """
+        Populate self.training_data for a single agent using the provided rollout trees.
+        """
+        pass
+
+    def set_trajectory_data(
+        self, roots: list[RolloutTreeRootNode], agent_ids: list[str]
+    ) -> None:
+        """
+        Convenience wrapper to ingest trajectory data for every training agent.
+        """
+        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:
+        """
+        Reset and rebuild the policy-gradient minibatches before iterating through agents.
+        """
+        self.policy_gradient_data = None
+        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:
+        """
+        Entry point for policy updates: prepare batches, compute gradients, and update parameters.
+        """
+        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()