larm/memory_generator/memgen_runner.py

import os
import logging
from torch.utils.data import DataLoader
from datasets import Dataset
from accelerate import Accelerator
from transformers import PreTrainedTokenizerBase
from trl import SFTTrainer, SFTConfig, GRPOConfig
from trl.models import unwrap_model_for_generation

from typing import Tuple, Dict, List
from tqdm import tqdm

from larm.task.base_runner import BaseRunner
from larm.common.registry import registry
from larm.common.config import Config
from larm.data.interactions.base_interaction import (
    InteractionConfig,
    InteractionManager,
    InteractionDataProto
)

from .memgen_model import LatentMemoryModel
from .trainer.weaver_grpo_trainer import WeaverGRPOTrainer
from .trainer.trigger_grpo_trainer import TriggerGRPOTrainer
from .utils import (
    fix_model_parameters,
    open_model_parameters,
    EvalConfig,
    StaticEvalRecorder,
    DynamicEvalRecorder
)
from PIL import Image
import torch


@registry.register_runner("latmem")
class LatentMemoryRunner(BaseRunner):

    def __init__(
        self,
        model: LatentMemoryModel,
        processing_class: PreTrainedTokenizerBase,
        datasets_dict: Dict,
        configs: Config,
        env_and_gens_dict: Dict
    ):
        super().__init__(
            model,
            processing_class,
            datasets_dict,
            configs,
            env_and_gens_dict
        )
        # parse configs
        self._parse_configs(configs.run_cfg)

        # initialize envs and generation managers
        dataset_config = configs.datasets_cfg[self.dataset_name]
        self.env = self.env_cls(dataset_config)

        # partition datasets
        self.weaver_train_dataset, self.trigger_train_dataset = self._parse_train_dataset(self.dataset_dict["train"])
        self.valid_dataset = self.dataset_dict["valid"]
        self.test_dataset = self.dataset_dict["test"]

        self.weaver_train_dataset = self._filter_dataset(self.weaver_train_dataset)
        self.trigger_train_dataset = self._filter_dataset(self.trigger_train_dataset)
        self.valid_dataset = self._filter_dataset(self.valid_dataset)

        # initialize generation manager
        self.generation_manager: InteractionManager = self.gen_cls(
            self.processing_class, self.model, self.interaction_config
        )

        # VIS: build multimodal features if present in dataset
        if self.weaver_train_dataset.column_names and any(k in self.weaver_train_dataset.column_names for k in ("image_path",)):
            self.weaver_train_dataset = self._prepare_mm_features(self.weaver_train_dataset)
        if self.valid_dataset.column_names and any(k in self.valid_dataset.column_names for k in ("image_path",)):
            self.valid_dataset = self._prepare_mm_features(self.valid_dataset)
        if self.test_dataset.column_names and any(k in self.test_dataset.column_names for k in ("image_path",)):
            self.test_dataset = self._prepare_mm_features(self.test_dataset)

    def _parse_train_dataset(self, train_dataset: Dataset) -> Tuple[Dataset, Dataset]:
        trigger_trainset_size = min(500, len(train_dataset))
        return train_dataset, train_dataset.select(range(trigger_trainset_size))

    def _filter_dataset(self, dataset: Dataset) -> Dataset:
        """Filter the dataset based on maximum sequence length.

        The maximum length depends on the training mode and method:
        - For Weaver SFT training: use `weaver_training_args.max_length`.
        - For Weaver GRPO training: use `weaver_training_args.max_prompt_length`.
        - For Trigger GRPO training: use `trigger_training_args.max_prompt_length`.

        Any sample exceeding the maximum length is filtered out.

        Args:
            dataset (Dataset): The input dataset to be filtered.

        Returns:
            Dataset: A filtered dataset containing only samples within the max length.
        """
        tokenizer = self.processing_class

        # Determine max length based on training mode
        max_len = 1024
        if self.train_weaver and self.train_weaver_method == "sft":
            max_len = self.weaver_training_args.max_length
        elif self.train_weaver and self.train_weaver_method == "grpo":
            max_len = self.weaver_training_args.max_prompt_length
        elif self.train_trigger and self.train_trigger_method == "grpo":
            max_len = self.trigger_training_args.max_prompt_length
        else:
            raise ValueError("Wrong training mode.")

        original_size = len(dataset)
        logging.info(f"[Filter] Starting filter with max_len={max_len}, dataset_size={original_size}")

        # Pre-extract tokenizer outside filter_func for better performance
        plain_tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class)

        # Function to filter out samples exceeding max length
        def filter_func(sample):
            # Check if sample is None or not a dict
            if sample is None or not isinstance(sample, dict):
                return False

            if "prompt" in sample and sample["prompt"] is not None:
                # Use pre-extracted tokenizer to avoid repeated getattr calls
                encoded = plain_tokenizer(sample["prompt"], add_special_tokens=True)
                return len(encoded["input_ids"]) < max_len
            elif "messages" in sample and sample["messages"] is not None:
                conversation = tokenizer.apply_chat_template(sample["messages"][:2], tokenize=True)
                return len(conversation) < max_len
            return True

        # Apply filtering with optimized settings
        logging.info(f"[Filter] Starting to apply filter function...")
        dataset = dataset.filter(
            filter_func,
            num_proc=None,  # 使用单进程（None 表示主进程，避免子进程序列化问题）
            load_from_cache_file=False,  # 禁用缓存避免卡住
            desc="Filter"
        )
        logging.info(f"[Filter] Filter function completed")

        filtered_size = len(dataset)
        logging.info(f"[Filter] Completed: {original_size} -> {filtered_size} samples (filtered out {original_size - filtered_size})")

        return dataset

    def _prepare_mm_features(self, dataset: Dataset) -> Dataset:
        processor = self.processing_class  # AutoProcessor for VL models
        tokenizer = getattr(processor, "tokenizer", processor)

        # Cache vision token IDs to avoid repeated encoding
        # For Qwen2.5-VL, check what vision tokens are actually used
        vision_start_ids = tokenizer.encode("<|vision_start|>", add_special_tokens=False)
        vision_end_ids = tokenizer.encode("<|vision_end|>", add_special_tokens=False)

        # Also try alternative vision token names for Qwen2.5-VL
        if len(vision_start_ids) == 0:
            # Try common alternatives
            for token in ["<|image_pad|>", "<|vision_pad|>", "<image>", "<img>"]:
                test_ids = tokenizer.encode(token, add_special_tokens=False)
                if len(test_ids) > 0:
                    vision_start_ids = test_ids
                    # logging.info(f"[MM_Features] Found vision start token: {token} -> {test_ids}")
                    break

        # logging.info(f"[MM_Features] Vision token IDs: start={vision_start_ids}, end={vision_end_ids}")
        # logging.info(f"[MM_Features] Tokenizer special tokens: {tokenizer.special_tokens_map}")

        def _encode(example: Dict) -> Dict:
            prompt = example.get("prompt")
            completion = example.get("completion")
            image_path = example.get("image_path")

            image = None
            # Debug: log only one image path candidate and its existence
            # try:
            #     _path_logged = getattr(_encode, "_path_logged_count", 0)
            # except Exception:
            #     _path_logged = 0
            # if _path_logged < 1:
            #     exists_str = os.path.exists(image_path) if image_path is not None else "N/A"
            #     logging.info(f"[MM_Features] example image_path: {image_path}, exists={exists_str}")
            #     try:
            #         _encode._path_logged_count = _path_logged + 1
            #     except Exception:
            #         pass
            if image_path is not None and os.path.exists(image_path):
                try:
                    image = Image.open(image_path).convert("RGB")
                except Exception:
                    image = None

            if image is not None:
                # Use Qwen2.5-VL's recommended message format
                # Build messages with image and text following official API
                messages = [
                    {
                        "role": "user",
                        "content": [
                            {"type": "image", "image": image},
                            {"type": "text", "text": prompt},
                        ],
                    }
                ]
                # Apply chat template first to insert vision tokens
                text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
                # Debug: log formatted text (once)
                # try:
                #     _template_logged = getattr(_encode, "_template_logged_count", 0)
                # except Exception:
                #     _template_logged = 0
                # if _template_logged < 1:
                #     logging.info(f"[MM_Features] Formatted text after apply_chat_template: {text[:200]}...")
                #     try:
                #         _encode._template_logged_count = _template_logged + 1
                #     except Exception:
                #         pass
                # Then process with image
                enc_prompt = processor(text=[text], images=[image], return_tensors="pt", padding=False)
                # Debug: confirm (log once) that we are encoding with image
                # try:
                #     _enc_logged = getattr(_encode, "_enc_with_img_logged_count", 0)
                # except Exception:
                #     _enc_logged = 0
                # if _enc_logged < 1:
                #     logging.info(f"[MM_Features] encoding with image: path={image_path}, size={getattr(image, 'size', None)}")
                #     logging.info(f"[MM_Features] enc_prompt keys: {list(enc_prompt.keys())}")
                #     # Check what processor returns
                #     for key, val in enc_prompt.items():
                #         if isinstance(val, torch.Tensor):
                #             logging.info(f"  {key}: Tensor{tuple(val.shape)}")
                #         else:
                #             logging.info(f"  {key}: {type(val)}")
                #     try:
                #         _encode._enc_with_img_logged_count = _enc_logged + 1
                #     except Exception:
                #         pass
                prompt_ids = enc_prompt["input_ids"][0]
                prompt_mask = enc_prompt["attention_mask"][0]
                pixel_values = enc_prompt.get("pixel_values")
                image_grid_thw = enc_prompt.get("image_grid_thw")
                if pixel_values is not None:
                    pixel_values = pixel_values[0]
                if image_grid_thw is not None:
                    image_grid_thw = image_grid_thw[0]
            else:
                enc_prompt = processor(text=[prompt], return_tensors="pt", padding=False)
                prompt_ids = enc_prompt["input_ids"][0]
                prompt_mask = enc_prompt["attention_mask"][0]
                pixel_values = None
                image_grid_thw = None

            tokenizer = getattr(processor, "tokenizer", processor)
            enc_comp = tokenizer(text=[completion], add_special_tokens=False, return_tensors="pt")
            comp_ids = enc_comp["input_ids"][0]

            input_ids = torch.cat([prompt_ids, comp_ids], dim=0)
            attention_mask = torch.cat([prompt_mask, torch.ones_like(comp_ids)], dim=0)
            labels = torch.cat([torch.full_like(prompt_ids, -100), comp_ids.clone()], dim=0)

            # For Qwen2.5-VL: image tokens are NOT explicitly in input_ids
            # Instead, the model uses image_grid_thw to locate image features
            # So image_token_mask may remain all False, which is correct
            image_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
            if image is not None:
                ids_list = prompt_ids.tolist()

                # Debug: log first example's token ids
                # try:
                #     _debug_logged = getattr(_encode, "_debug_token_ids_logged", 0)
                # except Exception:
                #     _debug_logged = 0
                # if _debug_logged < 1:
                #     logging.info(f"[MM_Features] Sample prompt_ids (first 50): {ids_list[:50]}")
                #     logging.info(f"[MM_Features] Looking for vision_start_ids: {vision_start_ids}, vision_end_ids: {vision_end_ids}")
                #     # Decode to show tokens
                #     decoded_tokens = [tokenizer.decode([tid]) for tid in ids_list[:50]]
                #     logging.info(f"[MM_Features] Decoded tokens (first 50): {decoded_tokens}")
                #     try:
                #         _encode._debug_token_ids_logged = 1
                #     except Exception:
                #         pass

                # For Qwen2.5-VL: check for <|vision_pad|> or <|image_pad|> tokens
                vision_pad_id = tokenizer.encode("<|vision_pad|>", add_special_tokens=False)
                image_pad_id = tokenizer.encode("<|image_pad|>", add_special_tokens=False)

                # if _debug_logged < 1:
                #     logging.info(f"[MM_Features] vision_pad_id: {vision_pad_id}, image_pad_id: {image_pad_id}")

                # Look for vision/image pad tokens in the sequence
                target_ids = []
                if len(vision_pad_id) > 0:
                    target_ids.append(vision_pad_id[0])
                if len(image_pad_id) > 0:
                    target_ids.append(image_pad_id[0])

                if len(target_ids) > 0:
                    # Mark all positions with vision_pad or image_pad tokens
                    for i, token_id in enumerate(ids_list):
                        if token_id in target_ids:
                            image_token_mask[i] = True

                    # if _debug_logged < 1:
                    #     num_marked = image_token_mask[:len(ids_list)].sum().item()
                    #     logging.info(f"[MM_Features] Marked {num_marked} image pad tokens (ids={target_ids})")

                # Also try vision_start/vision_end markers
                def find_subseq(seq, sub):
                    n, m = len(seq), len(sub)
                    for i in range(0, n - m + 1):
                        if seq[i:i+m] == sub:
                            return i
                    return -1

                s_idx = find_subseq(ids_list, vision_start_ids) if len(vision_start_ids) > 0 else -1
                e_idx = find_subseq(ids_list, vision_end_ids) if len(vision_end_ids) > 0 else -1

                # if _debug_logged < 1:
                #     logging.info(f"[MM_Features] Found indices: start_idx={s_idx}, end_idx={e_idx}")

                if s_idx != -1 and e_idx != -1 and e_idx >= s_idx:
                    image_token_mask[s_idx:e_idx+len(vision_end_ids)] = True
                    # if _debug_logged < 1:
                    #     num_marked = image_token_mask.sum().item()
                    #     logging.info(f"[MM_Features] Set image_token_mask[{s_idx}:{e_idx+len(vision_end_ids)}] = True (total marked: {num_marked})")

                # Final check
                # if _debug_logged < 1:
                #     total_marked = image_token_mask.sum().item()
                #     if total_marked == 0:
                #         logging.warning(f"[MM_Features] Image present but no vision tokens found in input_ids! This may be expected for Qwen2.5-VL.")

            out = {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "labels": labels,
                "image_token_mask": image_token_mask,
            }
            if pixel_values is not None:
                out["pixel_values"] = pixel_values
            if image_grid_thw is not None:
                out["image_grid_thw"] = image_grid_thw
            return out

        # Disable multiprocessing to avoid subprocess crashes with PIL/Processor
        # num_proc=None means no multiprocessing (runs in main process)
        # logging.info(f"[MM_Features] Processing {len(dataset)} samples with multimodal features...")
        # Safely get columns to remove, checking if column_names exists
        columns_to_remove = []
        if dataset.column_names:
            columns_to_remove = [c for c in dataset.column_names if c not in ("prompt", "completion", "solution", "image_path")]

        dataset = dataset.map(
            _encode,
            remove_columns=columns_to_remove,
            num_proc=None,  # Disable multiprocessing to avoid crashes
            load_from_cache_file=False,  # Disable cache due to closure serialization issues
            desc="Preparing multimodal features"
        )
        # logging.info(f"[MM_Features] Completed processing {len(dataset)} samples.")

        # Log distribution statistics of image_token_mask after processing
        try:
            sample_size = min(1000, len(dataset))
            counted = 0
            num_with_image_tokens = 0
            num_with_pixel_values = 0
            total_true_tokens = 0
            min_true_tokens = None
            max_true_tokens = None

            for i in range(sample_size):
                ex = dataset[i]
                mask = ex.get("image_token_mask", None)
                if mask is None:
                    continue

                # Check if this sample has actual pixel_values (real image data)
                pixel_values = ex.get("pixel_values", None)
                if pixel_values is not None:
                    num_with_pixel_values += 1

                # Support torch.Tensor, list, or numpy-like
                try:
                    if isinstance(mask, torch.Tensor):
                        cnt = int(mask.sum().item())
                    else:
                        cnt = int(sum(mask))
                except Exception:
                    try:
                        cnt = int(mask.sum())
                    except Exception:
                        cnt = 0

                counted += 1
                total_true_tokens += cnt
                if min_true_tokens is None or cnt < min_true_tokens:
                    min_true_tokens = cnt
                if max_true_tokens is None or cnt > max_true_tokens:
                    max_true_tokens = cnt
                if cnt > 0:
                    num_with_image_tokens += 1

            if counted > 0:
                mean_true_tokens = total_true_tokens / counted
                ratio_with_image = num_with_image_tokens / counted
                ratio_with_pixels = num_with_pixel_values / counted
                logging.info(
                    f"[MM_Features] Statistics on {counted} samples:"
                )
                logging.info(
                    f"  - Samples with pixel_values (actual images): {num_with_pixel_values}/{counted} ({ratio_with_pixels:.1%})"
                )
                logging.info(
                    f"  - Samples with vision tokens in input_ids: {num_with_image_tokens}/{counted} ({ratio_with_image:.1%})"
                )
                logging.info(
                    f"  - Vision tokens per sample: min/mean/max = {min_true_tokens}/{mean_true_tokens:.2f}/{max_true_tokens}"
                )
                if num_with_pixel_values > 0 and num_with_image_tokens == 0:
                    logging.warning(
                        f"[MM_Features] WARNING: {num_with_pixel_values} samples have images (pixel_values) "
                        f"but NO vision tokens found in input_ids. This is EXPECTED for Qwen2.5-VL, which uses "
                        f"image_grid_thw instead of explicit vision tokens in input_ids."
                    )
            else:
                logging.info("[MM_Features] image_token_mask: no samples counted (mask missing)")
        except Exception as e:
            logging.warning(f"[MM_Features] Failed to summarize image_token_mask distribution: {e}")

        # # Print a complete example after processing
        # try:
        #     if len(dataset) > 0:
        #         example = dataset[0]
        #         logging.info("=" * 80)
        #         logging.info("[MM_Features] COMPLETE EXAMPLE AFTER PROCESSING:")
        #         logging.info("=" * 80)

        #         # Print each field
        #         for key, value in example.items():
        #             if key == "pixel_values":
        #                 if value is not None:
        #                     if isinstance(value, torch.Tensor):
        #                         logging.info(f"  {key}: Tensor{tuple(value.shape)} dtype={value.dtype}")
        #                     else:
        #                         logging.info(f"  {key}: {type(value)}")
        #                 else:
        #                     logging.info(f"  {key}: None")

        #             elif key == "image_grid_thw":
        #                 if value is not None:
        #                     if isinstance(value, torch.Tensor):
        #                         logging.info(f"  {key}: Tensor{tuple(value.shape)} = {value.tolist()}")
        #                     else:
        #                         logging.info(f"  {key}: {value}")
        #                 else:
        #                     logging.info(f"  {key}: None")

        #             elif key == "input_ids":
        #                 if isinstance(value, torch.Tensor):
        #                     logging.info(f"  {key}: Tensor{tuple(value.shape)}")
        #                     # Show first 50 token IDs
        #                     logging.info(f"    Token IDs (first 50): {value.tolist()[:50]}")
        #                     # Decode each token individually to show special tokens
        #                     try:
        #                         tokens_display = []
        #                         for i in range(min(50, len(value))):
        #                             token_id = value[i].item()
        #                             token_str = tokenizer.decode([token_id], skip_special_tokens=False)
        #                             # Show token with its ID
        #                             tokens_display.append(f"{token_str}[{token_id}]")
        #                         logging.info(f"    Tokens (first 50): {' '.join(tokens_display)}")
        #                     except Exception as e:
        #                         logging.warning(f"    Failed to decode tokens: {e}")
        #                 else:
        #                     logging.info(f"  {key}: {value[:50]}...")

        #             elif key == "labels":
        #                 if isinstance(value, torch.Tensor):
        #                     # Count non-ignored labels
        #                     num_labels = (value != -100).sum().item()
        #                     logging.info(f"  {key}: Tensor{tuple(value.shape)}, non-ignored labels={num_labels}")
        #                     # Show where labels start (first non -100 position)
        #                     non_ignore_indices = (value != -100).nonzero(as_tuple=True)[0]
        #                     if len(non_ignore_indices) > 0:
        #                         first_label_idx = non_ignore_indices[0].item()
        #                         logging.info(f"    Labels start at position {first_label_idx}")
        #                         # Decode the label tokens (non -100 parts)
        #                         try:
        #                             label_tokens = value[value != -100][:30]  # First 30 label tokens
        #                             tokens_display = []
        #                             for token_id in label_tokens:
        #                                 token_str = tokenizer.decode([token_id.item()], skip_special_tokens=False)
        #                                 tokens_display.append(f"{token_str}[{token_id.item()}]")
        #                             logging.info(f"    Label tokens (first 30): {' '.join(tokens_display)}")
        #                         except Exception as e:
        #                             logging.warning(f"    Failed to decode label tokens: {e}")
        #                 else:
        #                     logging.info(f"  {key}: {value[:50]}...")

        #             elif key == "attention_mask":
        #                 if isinstance(value, torch.Tensor):
        #                     num_valid = int(value.sum().item())
        #                     logging.info(f"  {key}: Tensor{tuple(value.shape)}, valid positions={num_valid}/{len(value)}")
        #                     logging.info(f"    First 50 values: {value.tolist()[:50]}")
        #                 else:
        #                     logging.info(f"  {key}: {value[:50]}...")

        #             elif key == "image_token_mask":
        #                 if isinstance(value, torch.Tensor):
        #                     num_image_tokens = int(value.sum().item())
        #                     logging.info(f"  {key}: Tensor{tuple(value.shape)}, image token positions={num_image_tokens}")
        #                     if num_image_tokens > 0:
        #                         # Show which positions are image tokens
        #                         image_positions = value.nonzero(as_tuple=True)[0].tolist()[:20]
        #                         logging.info(f"    Image token positions (first 20): {image_positions}")
        #                     else:
        #                         logging.info(f"    No image tokens found in input_ids (expected for Qwen2.5-VL)")
        #                 else:
        #                     logging.info(f"  {key}: {value[:50]}...")

        #             else:
        #                 # Other fields (prompt, completion, solution, image_path)
        #                 if isinstance(value, str):
        #                     logging.info(f"  {key}: '{value[:200]}{'...' if len(value) > 200 else ''}'")
        #                 else:
        #                     logging.info(f"  {key}: {value}")

        #         logging.info("=" * 80)
        # except Exception as e:
        #     logging.warning(f"[MM_Features] Failed to print example: {e}")

        return dataset

    # ===== train weaver =====
    def _create_weaver_trainer(self):

        # SFT Trainer
        if self.train_weaver_method == "sft":
            # JSONL log path under save_dir
            weaver_trainer = SFTTrainer(
                model=self.model,
                args=self.weaver_training_args,
                train_dataset=self.weaver_train_dataset,
                eval_dataset=self.valid_dataset,
                processing_class=self.processing_class,
            )

        # GRPO Trainer
        elif self.train_weaver_method == 'grpo':

            reward_funcs = []     # get reward funcs for GRPO Trainer
            for reward_name in self.weaver_reward_names:
                reward_funcs.append(self.env_cls.get_reward_func(reward_name))

            weaver_trainer = WeaverGRPOTrainer(
                model=self.model,
                reward_funcs=reward_funcs,
                args=self.weaver_training_args,
                train_dataset=self.weaver_train_dataset,
                eval_dataset=self.valid_dataset,
                processing_class=self.processing_class,
                env_class=self.env_cls,
                env_main_config=self.configs.datasets_cfg[self.dataset_name],
                generation_manager=self.generation_manager
            )
        else:
            raise ValueError("Unsupported weaver training method.")

        return weaver_trainer

    def _train_weaver(self):

        # fix trigger parameters
        fix_model_parameters(self.model.trigger)

        # train weaver
        weaver_trainer = self._create_weaver_trainer()
        weaver_trainer.train()
        weaver_trainer.save_model()   # save the best model

        # remove checkpoints and save weaver
        output_dir = weaver_trainer.args.output_dir
        self._remove_trainer_ckpts(output_dir)

        # open trigger parameters
        open_model_parameters(self.model.trigger)

    # ===== train trigger =====

    def _create_trigger_trainer(self):

        # get reward funcs for RL Trainer
        reward_funcs = []
        for reward_name in self.trigger_reward_names:
            reward_funcs.append(self.env_cls.get_reward_func(reward_name))

        # build trainer
        trigger_trainer = TriggerGRPOTrainer(
            model=self.model,
            processing_class=self.processing_class,
            train_dataset=self.trigger_train_dataset,
            eval_dataset=self.valid_dataset,
            reward_funcs=reward_funcs,
            args=self.trigger_training_args
        )

        return trigger_trainer

    def _train_trigger(self):

        # fix weaver parameters
        fix_model_parameters(self.model.weaver)

        # train trigger
        trigger_trainer = self._create_trigger_trainer()
        trigger_trainer.train()
        trigger_trainer.save_model()     # save the best model

        # remove checkpoints and save weaver
        output_dir = trigger_trainer.args.output_dir
        self._remove_trainer_ckpts(output_dir)

        # open trigger parameters
        open_model_parameters(self.model.weaver)

    # ===== train weaver/trigger =====
    def train(self):
        # train weaver
        if self.train_weaver:
            self._train_weaver()

        # train trigger
        if self.train_trigger:
            self._train_trigger()

    # ===== evaluate =====
    def _static_evaluate(self):

        accelerator = Accelerator()
        writer = self._create_tensorboard(mode="evaluate")

        batch_size = self.eval_config.batch_size
        output_dir = self.eval_config.output_dir
        generation_config = self.eval_config.generation_config
        _actual_tokenizer = getattr(self.processing_class, 'tokenizer', self.processing_class)
        generation_config.eos_token_id = _actual_tokenizer.eos_token_id

        # prepare dataset and dataloader
        test_dataloader = accelerator.prepare(DataLoader(
            dataset=self.test_dataset,
            batch_size=batch_size,
            shuffle=False,
            collate_fn=lambda batch: batch  # use the identity function
        ))

        # prepare model
        model_wrapped = accelerator.prepare_model(model=self.model, evaluation_mode=True)
        model_wrapped.eval()

        # construct eval recorder
        test_funcs = [self.env_cls.get_reward_func("accuracy")]
        save_file = os.path.join(output_dir, "answer.json")
        recorder = StaticEvalRecorder(compute_metrics=test_funcs, writer=writer, log_file=save_file)

        # batch generation
        for test_batch in tqdm(test_dataloader):
            with unwrap_model_for_generation(
                model_wrapped, accelerator
            ) as unwrapped_model:
                # construct InteractionDataProto object
                prompts = [x["prompt"] for x in test_batch]
                prompt_inputs = self.processing_class(
                    text=prompts, return_tensors="pt", padding=True, padding_side="left", add_special_tokens=True
                )
                prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]
                gen_batch = InteractionDataProto()
                gen_batch.batch["input_ids"] = prompt_ids.to(accelerator.device)
                gen_batch.batch["attention_mask"] = prompt_mask.to(accelerator.device)
                gen_batch.no_tensor_batch["initial_prompts"] = [x["prompt"] for x in test_batch]

                # generation manager
                self.generation_manager.actor_rollout_wg = unwrapped_model
                gen_output = self.generation_manager.run_agent_loop(gen_batch)

                # postprocess: 由 generation manager 保证 completion ids 的正确性
                completion_ids = gen_output.batch["responses"]
                completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True)

            # 转为 seperated examples
            recorder.record_batch(completions, test_batch)
        recorder.finalize()
        writer.close()

    def _dynamic_evaluate(self):

        def _set_batch_envs(batch: List) -> Tuple[List[str], List[str], List]:  # batch set envs
            system_prompts, init_user_prompts, envs = [], [], []
            for task_config in batch:
                env = self.env_cls(self.configs.datasets_cfg[self.dataset_name])
                system_prompt, init_user_prompt = env.set_env(task_config)

                system_prompts.append(system_prompt)
                init_user_prompts.append(init_user_prompt)
                envs.append(env)

            return system_prompts, init_user_prompts, envs

        def _build_data_proto(
            system_prompts: List[str], init_user_prompts: List[str], envs: List
        ) -> InteractionDataProto:
            messages = []
            for system_prmopt, init_user_prompt in zip(system_prompts, init_user_prompts):
                system_message = {"role": "system", "content": system_prmopt}
                user_message = {"role": "user", "content": init_user_prompt}
                init_messages = [system_message, user_message]
                messages.append(init_messages)

            data_proto = InteractionDataProto()
            data_proto.no_tensor_batch["init_prompts"] = messages
            data_proto.no_tensor_batch["envs"] = envs

            return data_proto

        # ===== body =====
        output_dir = self.eval_config.output_dir

        accelerator = Accelerator()
        writer = self._create_tensorboard(mode="evaluate")
        save_file = os.path.join(output_dir, "conversations.txt")
        recorder = DynamicEvalRecorder(writer=writer, log_file=save_file)

        batch_size = self.eval_config.batch_size
        generation_config = self.eval_config.generation_config
        _actual_tokenizer = getattr(self.processing_class, 'tokenizer', self.processing_class)
        generation_config.eos_token_id = _actual_tokenizer.eos_token_id

        # prepare dataset and dataloader
        test_dataloader = accelerator.prepare(DataLoader(
            dataset=self.test_dataset,
            batch_size=batch_size,
            shuffle=False,
            collate_fn=lambda batch: batch  # use the identity function
        ))

        # prepare model
        model_wrapped = accelerator.prepare_model(model=self.model, evaluation_mode=True)
        model_wrapped.eval()

        # batch generate
        for step, test_batch in tqdm(enumerate(test_dataloader)):
            with unwrap_model_for_generation(
                model_wrapped, accelerator
            ) as unwrapped_model:
                system_prompts, init_user_prompts, envs = _set_batch_envs(test_batch)
                input_data_proto = _build_data_proto(system_prompts, init_user_prompts, envs)

                self.generation_manager.actor_rollout_wg = unwrapped_model
                outputs: InteractionDataProto = self.generation_manager.run_agent_loop(input_data_proto)

                inter_histories = outputs.no_tensor_batch["inter_histories"]
                inter_context = self.processing_class.apply_chat_template(inter_histories, tokenize=False)

            # batch record
            rewards = []
            for env in input_data_proto.no_tensor_batch["envs"]:
                reward = env.feedback()
                rewards.append(reward)

            recorder.record_batch(inter_context, rewards)

        recorder.finalize()
        writer.close()

    # ===== runner configs =====

    def _parse_configs(self, configs):
        """build configs
        - weaver training config
        - trigger training config
        - interaction config
        - evaluatoin config
        """
        self.save_dir = configs.get("save_dir")
        use_tensorboard = configs.get("use_wandb")

        # weaver configs
        self.train_weaver = configs.get("train_weaver", True)
        if self.train_weaver:
            self.train_weaver_method = configs.get("train_weaver_method", "sft")
            weaver_save_dir = os.path.join(self.save_dir, "weaver")
            weaver_config = configs.get("weaver", {})

            # train weaver with sft
            if self.train_weaver_method == "sft":
                sft_config = weaver_config.get("sft", {})
                weaver_args_dict = self._parse_common_training_args(sft_config, weaver_save_dir, use_tensorboard)
                self.weaver_training_args = SFTConfig(**weaver_args_dict)

            # train weaver with grpo
            elif self.train_weaver_method == "grpo":
                grpo_config = weaver_config.get("grpo", {})
                weaver_args_dict = self._parse_common_training_args(grpo_config, weaver_save_dir, use_tensorboard, is_grpo=True)
                self.weaver_reward_names = weaver_args_dict.pop("reward_names")

                self.weaver_training_args = GRPOConfig(**weaver_args_dict)

            else:
                raise ValueError("Unsupported weaver training mode")

        # Trigger configs
        self.train_trigger = configs.get("train_trigger", False)
        if self.train_trigger:
            self.train_trigger_method = configs.get("train_trigger_method", "grpo")
            trigger_save_dir = os.path.join(self.save_dir, "trigger")
            trigger_config = configs.get("trigger", {})

            if self.train_trigger_method == "grpo":
                grpo_config = trigger_config.get("grpo", {})
                trigger_args_dict = self._parse_common_training_args(grpo_config, trigger_save_dir, use_tensorboard, is_grpo=True)
                self.trigger_reward_names = trigger_args_dict.pop("reward_names")

                self.trigger_training_args = GRPOConfig(**trigger_args_dict)
            else:
                raise ValueError("Unsupported weaver training mode")

        # Interaction configs
        generation_configs = configs.get("generation", {})
        self.interaction_config = InteractionConfig(
            max_turns=generation_configs.get("max_turns", 30),
            max_start_length=generation_configs.get("max_start_length", 1024),
            max_prompt_length=generation_configs.get("max_prompt_length", 4096),
            max_response_length=generation_configs.get("max_response_length", 512),
            max_obs_length=generation_configs.get("max_obs_length", 512),
            do_sample=generation_configs.get("do_sample", False),
            temperature=generation_configs.get("temperature", 1.0)
        )

        # Evaluation configs
        eval_dir = os.path.join(self.save_dir, "evaluate")
        eval_batch_size = generation_configs.get("eval_batch_size", 32)
        self.eval_config = EvalConfig(
            output_dir=eval_dir, batch_size=eval_batch_size, generation_config=self.interaction_config
        )

        # Align GRPO generation configuration with the interaction configuration:
        # All generation-related settings are controlled by the interaction config.
        if (self.train_weaver and self.train_weaver_method == "grpo"):
            self.weaver_training_args.max_prompt_length = self.interaction_config.max_start_length
            self.weaver_training_args.max_completion_length = self.interaction_config.max_response_length
            self.weaver_training_args.temperature = self.interaction_config.temperature
        elif (self.train_trigger and self.train_trigger_method == "grpo"):
            self.trigger_training_args.max_prompt_length = self.interaction_config.max_start_length
            self.trigger_training_args.max_completion_length = self.interaction_config.max_response_length
            self.trigger_training_args.temperature = self.interaction_config.temperature

    def _parse_common_training_args(self, config_dict, output_dir, use_tensorboard, is_grpo=False):
        batch_size = config_dict.get("batch_size", 4)
        max_epochs = config_dict.get("max_epochs", 2)
        grad_accum_steps = config_dict.get("grad_accum_steps", 1)

        optim = config_dict.get("optim", "adamw_torch")
        lr = config_dict.get("lr", 1e-5)
        scheduler = config_dict.get("schedular", "cosine")
        warmup_ratio = config_dict.get("warmup_ratio", 0.1)

        logging_strategy = config_dict.get("logging_strategy", "steps")
        logging_steps = config_dict.get("logging_steps", 1) if logging_strategy == "steps" else None

        eval_strategy = config_dict.get("eval_strategy", "steps")
        eval_steps = config_dict.get("eval_steps", 200) if eval_strategy == "steps" else None

        save_strategy = config_dict.get("save_strategy", "steps")
        save_steps = config_dict.get("save_steps", 200) if save_strategy == "steps" else None

        # common args dict
        args_dict = {
            "output_dir": output_dir,
            "per_device_train_batch_size": batch_size,
            "per_device_eval_batch_size": batch_size,
            "num_train_epochs": max_epochs,
            "gradient_accumulation_steps": grad_accum_steps,
            "optim": optim,
            "learning_rate": lr,
            "lr_scheduler_type": scheduler,
            "warmup_ratio": warmup_ratio,
            "logging_strategy": logging_strategy,
            "logging_steps": logging_steps,
            "save_strategy": save_strategy,
            "save_steps": save_steps,
            "eval_strategy": eval_strategy,
            "eval_steps": eval_steps,
            "report_to": ["tensorboard"] if use_tensorboard else [],   # <-- set to tensorboard
            "remove_unused_columns": False,
            "load_best_model_at_end": True,
            "bf16": True,
        }

        # add grpo specific args
        if is_grpo:
            args_dict.update({
                "num_generations": config_dict.get("num_generations", 16),
                "num_iterations": config_dict.get("num_iterations", 1),
                "beta": config_dict.get("beta", 0.0),
                "loss_type": config_dict.get("loss_type", "grpo"),
                "max_prompt_length": config_dict.get("max_prompt_length", 1024),
                "max_completion_length": config_dict.get("max_completion_length", 512),
            })

            rewards = config_dict.get("reward_funcs", [])
            reward_weights = [r["weight"] for r in rewards]
            reward_names = [r["name"] for r in rewards]

            args_dict.update({
                "reward_weights": reward_weights,
                "reward_names": reward_names
            })
        # add sft specific args
        else:
            args_dict.update({
                "max_length": config_dict.get("max_length", 1024),
                "assistant_only_loss": config_dict.get("assistant_only_loss", True)
            })

        return args_dict