From 9c2a5811b6061a65ead3718f72ad46f07e703b9b Mon Sep 17 00:00:00 2001 From: N!no Date: Sat, 4 Jul 2026 12:47:13 -0400 Subject: [PATCH 1/4] add centralized exploration --- comlrl/trainers/preference/iterative.py | 477 ++++++++++++++++++++++-- 1 file changed, 451 insertions(+), 26 deletions(-) diff --git a/comlrl/trainers/preference/iterative.py b/comlrl/trainers/preference/iterative.py index ffe5478..1af8990 100644 --- a/comlrl/trainers/preference/iterative.py +++ b/comlrl/trainers/preference/iterative.py @@ -3,6 +3,7 @@ import json import os import random +import re from dataclasses import dataclass from typing import Any, Dict, List, Optional, Sequence, Tuple, Union from urllib import error as urlerror @@ -54,6 +55,17 @@ def _normalize_comparator_policy(policy: Optional[str]) -> str: raise ValueError("comparator_policy must be one of: current, model, history, api.") +def _normalize_comparator_generation_mode(mode: Optional[str]) -> str: + value = str(mode or "decentralized").strip().lower() + if value in {"decentralized", "decentralised", "multi_agent", "multi-agent"}: + return "decentralized" + if value in {"centralized", "centralised", "single_agent", "single-agent"}: + return "centralized" + raise ValueError( + "comparator_generation_mode must be one of: decentralized, centralized." + ) + + def _normalize_preference_replay_mode(mode: Optional[str]) -> str: value = str(mode or "current").strip().lower() if value in {"current", "latest", "new"}: @@ -86,6 +98,24 @@ def _validate_iterative_config(args: Any) -> None: ) _validate_preference_replay_args(args) args.comparator_policy = _normalize_comparator_policy(args.comparator_policy) + args.comparator_generation_mode = _normalize_comparator_generation_mode( + getattr(args, "comparator_generation_mode", "decentralized") + ) + if args.comparator_generation_mode == "centralized": + if int(args.num_agents) != 2: + raise ValueError( + "comparator_generation_mode='centralized' currently requires " + "num_agents=2." + ) + if int(args.comparator_centralized_agent_index) < 0: + raise ValueError("comparator_centralized_agent_index must be >= 0.") + if int(args.comparator_centralized_agent_index) >= int(args.num_agents): + raise ValueError( + "comparator_centralized_agent_index must be smaller than num_agents." + ) + args.comparator_centralized_agent_index = int( + args.comparator_centralized_agent_index + ) if args.comparator_num_candidates is not None: if int(args.comparator_num_candidates) < 1: raise ValueError("comparator_num_candidates must be >= 1 or null.") @@ -229,8 +259,11 @@ class MADPOIterConfig(MADPOConfig): preference_replay_sample_size: Optional[int] = None preference_replay_dir: Optional[str] = None log_reward_distribution: bool = False + log_comparator_generations: bool = False policy_checkpoint_dir: Optional[str] = None comparator_policy: str = "current" + comparator_generation_mode: str = "decentralized" + comparator_centralized_agent_index: int = 0 comparator_model_name: Optional[str] = None comparator_agents: Optional[Sequence[str]] = None comparator_devices: Optional[Union[str, Sequence[str]]] = None @@ -273,9 +306,12 @@ class MARLHFIterConfig(MARLHFConfig): preference_replay_sample_size: Optional[int] = None preference_replay_dir: Optional[str] = None log_reward_distribution: bool = False + log_comparator_generations: bool = False preference_scoring_reward: str = "task" policy_checkpoint_dir: Optional[str] = None comparator_policy: str = "current" + comparator_generation_mode: str = "decentralized" + comparator_centralized_agent_index: int = 0 comparator_model_name: Optional[str] = None comparator_agents: Optional[Sequence[str]] = None comparator_devices: Optional[Union[str, Sequence[str]]] = None @@ -1151,6 +1187,62 @@ def _reward_distribution_dir(self) -> str: self._reward_distribution_dir_path = path return path + def _comparator_generation_dir(self) -> str: + cached = getattr(self, "_comparator_generation_dir_path", None) + if cached: + return cached + + path = None + if isinstance(self.wandb_config, dict): + output_dir = self.wandb_config.get("output_dir") + if output_dir: + path = os.path.join(str(output_dir), "comparator_generations") + else: + sections = self.wandb_config.get("config_sections") or {} + output_section = ( + sections.get("output") if isinstance(sections, dict) else {} + ) + base_dir = None + if isinstance(output_section, dict): + base_dir = output_section.get("base_dir") + base_dir = base_dir or self.wandb_config.get("dir") + if base_dir: + job_id = os.environ.get("SLURM_JOB_ID") + path = ( + os.path.join( + str(base_dir), f"job_{job_id}", "comparator_generations" + ) + if job_id + else os.path.join(str(base_dir), "comparator_generations") + ) + if not path: + path = os.path.join(os.getcwd(), "comparator_generations") + + path = os.path.abspath(str(path)) + os.makedirs(path, exist_ok=True) + self._comparator_generation_dir_path = path + return path + + def _write_comparator_generation_record( + self, + iteration_idx: int, + record: Dict[str, Any], + ) -> None: + if not getattr(self.args, "log_comparator_generations", False): + return + path = os.path.join( + self._comparator_generation_dir(), + f"iteration_{iteration_idx + 1:04d}.jsonl", + ) + payload = { + "iteration": int(iteration_idx + 1), + "comparator_policy": self.args.comparator_policy, + "comparator_generation_mode": self.args.comparator_generation_mode, + **record, + } + with open(path, "a", encoding="utf-8") as handle: + handle.write(json.dumps(self._jsonable(payload), ensure_ascii=False) + "\n") + def _write_iteration_preference_pairs( self, iteration_idx: int, @@ -1454,32 +1546,12 @@ def _generate_preference_pairs_for_item( num_candidates=current_candidates, **kwargs, ) - if self.args.comparator_policy == "api": - comparator_outputs = self._generate_api_outputs_for_item( - batch_item, - num_candidates=comparator_candidates, - ) - elif self.args.comparator_policy == "current": - comparator_outputs = self._generate_policy_outputs_for_item( - self.agents, - batch_item, - num_candidates=comparator_candidates, - **kwargs, - ) - elif self.args.comparator_policy == "history": - comparator_outputs = self._generate_history_policy_outputs_for_item( - batch_item, - iteration_idx=iteration_idx, - num_candidates=comparator_candidates, - **kwargs, - ) - else: - comparator_outputs = self._generate_policy_outputs_for_item( - self._get_comparator_agents(), - batch_item, - num_candidates=comparator_candidates, - **kwargs, - ) + comparator_outputs = self._generate_comparator_outputs_for_item( + batch_item, + iteration_idx=iteration_idx, + num_candidates=comparator_candidates, + **kwargs, + ) current_completions = [ current_outputs[i]["completions"][0] for i in range(self.num_agents) @@ -1674,6 +1746,72 @@ def _preference_scoring_raw_and_processed_rewards( ) -> Optional[Tuple[List[float], List[float]]]: return None + def _generate_comparator_outputs_for_item( + self, + batch_item: Dict[str, Any], + *, + iteration_idx: int, + num_candidates: int, + **kwargs, + ) -> List[Dict[str, Any]]: + if self.args.comparator_generation_mode == "centralized": + return self._generate_centralized_comparator_outputs_for_item( + batch_item, + iteration_idx=iteration_idx, + num_candidates=num_candidates, + **kwargs, + ) + + if self.args.comparator_policy == "api": + outputs = self._generate_api_outputs_for_item( + batch_item, + num_candidates=num_candidates, + ) + self._log_decentralized_comparator_outputs( + outputs, + batch_item=batch_item, + iteration_idx=iteration_idx, + ) + return outputs + if self.args.comparator_policy == "current": + outputs = self._generate_policy_outputs_for_item( + self.agents, + batch_item, + num_candidates=num_candidates, + **kwargs, + ) + self._log_decentralized_comparator_outputs( + outputs, + batch_item=batch_item, + iteration_idx=iteration_idx, + ) + return outputs + if self.args.comparator_policy == "history": + outputs = self._generate_history_policy_outputs_for_item( + batch_item, + iteration_idx=iteration_idx, + num_candidates=num_candidates, + **kwargs, + ) + self._log_decentralized_comparator_outputs( + outputs, + batch_item=batch_item, + iteration_idx=iteration_idx, + ) + return outputs + outputs = self._generate_policy_outputs_for_item( + self._get_comparator_agents(), + batch_item, + num_candidates=num_candidates, + **kwargs, + ) + self._log_decentralized_comparator_outputs( + outputs, + batch_item=batch_item, + iteration_idx=iteration_idx, + ) + return outputs + def _generate_policy_outputs_for_item( self, policy_agents: Sequence[Any], @@ -1694,6 +1832,293 @@ def _generate_agent(agent_idx: int) -> Dict[str, Any]: return self._run_agent_tasks(_generate_agent) + def _generate_centralized_comparator_outputs_for_item( + self, + batch_item: Dict[str, Any], + *, + iteration_idx: int, + num_candidates: int, + **kwargs, + ) -> List[Dict[str, Any]]: + prompt = self._build_centralized_comparator_prompt(batch_item) + if self.args.comparator_policy == "api": + completions = self._call_comparator_api( + prompt=prompt, + agent_idx=self._centralized_comparator_agent_index(), + batch_item=batch_item, + num_candidates=num_candidates, + ) + return self._split_centralized_comparator_outputs( + completions, + batch_item=batch_item, + prompt=prompt, + iteration_idx=iteration_idx, + ) + if self.args.comparator_policy == "current": + return self._generate_centralized_policy_outputs_for_item( + self.agents, + batch_item, + prompt=prompt, + iteration_idx=iteration_idx, + num_candidates=num_candidates, + **kwargs, + ) + if self.args.comparator_policy == "history": + checkpoint_dir = self._history_policy_checkpoint_path(iteration_idx) + comparator_agents = self._load_policy_checkpoint_agents(checkpoint_dir) + try: + return self._generate_centralized_policy_outputs_for_item( + comparator_agents, + batch_item, + prompt=prompt, + iteration_idx=iteration_idx, + num_candidates=num_candidates, + **kwargs, + ) + finally: + del comparator_agents + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return self._generate_centralized_policy_outputs_for_item( + self._get_comparator_agents(), + batch_item, + prompt=prompt, + iteration_idx=iteration_idx, + num_candidates=num_candidates, + **kwargs, + ) + + def _generate_centralized_policy_outputs_for_item( + self, + policy_agents: Sequence[Any], + batch_item: Dict[str, Any], + *, + prompt: str, + iteration_idx: int, + num_candidates: int, + **kwargs, + ) -> List[Dict[str, Any]]: + agent_idx = self._centralized_comparator_agent_index() + generation_output = self._generate_completions( + policy_agents[agent_idx], + [batch_item], + agent_idx=agent_idx, + num_return_sequences=num_candidates, + max_new_tokens=self.args.max_new_tokens, + prompts_override=[prompt], + **kwargs, + ) + completions = generation_output["completions"][0] + return self._split_centralized_comparator_outputs( + completions, + batch_item=batch_item, + prompt=prompt, + iteration_idx=iteration_idx, + ) + + def _centralized_comparator_agent_index(self) -> int: + return int(getattr(self.args, "comparator_centralized_agent_index", 0)) + + def _build_centralized_comparator_prompt(self, batch_item: Dict[str, Any]) -> str: + if self.num_agents != 2: + raise ValueError( + "Centralized comparator generation currently supports exactly 2 agents." + ) + auxiliary_prompt = self.formatters[0](batch_item) + main_prompt = self.formatters[1](batch_item) + return f"""You are acting as one centralized coordinator for two code-generation agents. + +Your job is to produce the exact outputs that two decentralized agents would submit: +one Auxiliary output and one Main output. The two outputs must work together. + +Auxiliary agent original prompt: +{auxiliary_prompt} + +Main agent original prompt: +{main_prompt} + +IMPORTANT INSTRUCTIONS: +- Return exactly two complete code snippets: one for Auxiliary and one for Main. +- The Auxiliary snippet must define the helper function aux(...). +- The Main snippet must define only the required main function from the problem. +- Do not include explanations, examples, tests, markdown fences, or extra text. +- Use the exact section tags below so the parser can split the two outputs. + + +def aux(...): + # your auxiliary code here + return result + + +
+def required_function(...): + # your main code here + return result +
+""" + + def _split_centralized_comparator_outputs( + self, + completions: Sequence[str], + *, + batch_item: Dict[str, Any], + prompt: str, + iteration_idx: int, + ) -> List[Dict[str, Any]]: + auxiliary_completions: List[str] = [] + main_completions: List[str] = [] + for candidate_idx, completion in enumerate(completions): + raw_completion = str(completion) + try: + auxiliary, main = self._parse_centralized_comparator_completion( + raw_completion + ) + except Exception as exc: + self._write_comparator_generation_record( + iteration_idx, + { + "candidate_index": int(candidate_idx), + "centralized_agent_index": ( + self._centralized_comparator_agent_index() + ), + "batch_item": self._jsonable(batch_item), + "prompt": prompt, + "raw_completion": raw_completion, + "parsed_auxiliary": None, + "parsed_main": None, + "parse_ok": False, + "parse_error": str(exc), + }, + ) + raise + auxiliary_completions.append(auxiliary) + main_completions.append(main) + self._write_comparator_generation_record( + iteration_idx, + { + "candidate_index": int(candidate_idx), + "centralized_agent_index": ( + self._centralized_comparator_agent_index() + ), + "batch_item": self._jsonable(batch_item), + "prompt": prompt, + "raw_completion": raw_completion, + "parsed_auxiliary": auxiliary, + "parsed_main": main, + "parse_ok": True, + "parse_error": None, + }, + ) + + return [ + { + "prompts": [prompt], + "batch_items": [batch_item], + "completions": [auxiliary_completions], + }, + { + "prompts": [prompt], + "batch_items": [batch_item], + "completions": [main_completions], + }, + ] + + def _parse_centralized_comparator_completion(self, text: str) -> Tuple[str, str]: + auxiliary = self._extract_tagged_section(text, "auxiliary") + main = self._extract_tagged_section(text, "main") + if auxiliary is not None and main is not None: + return ( + self._clean_centralized_section(auxiliary), + self._clean_centralized_section(main), + ) + + auxiliary, main = self._extract_labeled_centralized_sections(text) + if auxiliary is None or main is None: + raise ValueError( + "Centralized comparator output could not be parsed. Expected " + "... and
...
sections." + ) + return ( + self._clean_centralized_section(auxiliary), + self._clean_centralized_section(main), + ) + + @staticmethod + def _extract_tagged_section(text: str, tag: str) -> Optional[str]: + pattern = rf"<\s*{tag}\s*>(.*?)<\s*/\s*{tag}\s*>" + match = re.search(pattern, text, flags=re.IGNORECASE | re.DOTALL) + if match: + return match.group(1) + return None + + @staticmethod + def _extract_labeled_centralized_sections( + text: str, + ) -> Tuple[Optional[str], Optional[str]]: + auxiliary_label = re.search( + r"(?im)^\s*(?:#{1,6}\s*)?(?:auxiliary|aux)\s*(?:code|output)?\s*:\s*$", + text, + ) + main_label = re.search( + r"(?im)^\s*(?:#{1,6}\s*)?main\s*(?:code|output)?\s*:\s*$", + text, + ) + if auxiliary_label is None or main_label is None: + return None, None + + if auxiliary_label.start() < main_label.start(): + auxiliary = text[auxiliary_label.end() : main_label.start()] + main = text[main_label.end() :] + else: + main = text[main_label.end() : auxiliary_label.start()] + auxiliary = text[auxiliary_label.end() :] + return auxiliary, main + + @staticmethod + def _clean_centralized_section(text: str) -> str: + value = text.strip() + fence_match = re.fullmatch( + r"```(?:python)?\s*(.*?)```", value, flags=re.IGNORECASE | re.DOTALL + ) + if fence_match: + value = fence_match.group(1).strip() + return value + + def _log_decentralized_comparator_outputs( + self, + outputs: Sequence[Dict[str, Any]], + *, + batch_item: Dict[str, Any], + iteration_idx: int, + ) -> None: + if not getattr(self.args, "log_comparator_generations", False): + return + for agent_idx, output in enumerate(outputs): + prompts = output.get("prompts") if isinstance(output, dict) else None + prompt = prompts[0] if isinstance(prompts, list) and prompts else None + completions = ( + output.get("completions") if isinstance(output, dict) else None + ) + candidate_completions = ( + completions[0] if isinstance(completions, list) and completions else [] + ) + for candidate_idx, completion in enumerate(candidate_completions): + self._write_comparator_generation_record( + iteration_idx, + { + "agent_idx": int(agent_idx), + "candidate_index": int(candidate_idx), + "batch_item": self._jsonable(batch_item), + "prompt": prompt, + "raw_completion": str(completion), + "parsed_auxiliary": None, + "parsed_main": None, + "parse_ok": True, + "parse_error": None, + }, + ) + def _generate_history_policy_outputs_for_item( self, batch_item: Dict[str, Any], From 621e45c0f4110137e19478790e93336ece65db0f Mon Sep 17 00:00:00 2001 From: N!no Date: Sat, 4 Jul 2026 13:17:12 -0400 Subject: [PATCH 2/4] Update iterative.py --- comlrl/trainers/preference/iterative.py | 145 +++++++++++++++++++++++- 1 file changed, 142 insertions(+), 3 deletions(-) diff --git a/comlrl/trainers/preference/iterative.py b/comlrl/trainers/preference/iterative.py index 1af8990..9f1cbf2 100644 --- a/comlrl/trainers/preference/iterative.py +++ b/comlrl/trainers/preference/iterative.py @@ -1972,7 +1972,8 @@ def _split_centralized_comparator_outputs( raw_completion = str(completion) try: auxiliary, main = self._parse_centralized_comparator_completion( - raw_completion + raw_completion, + batch_item=batch_item, ) except Exception as exc: self._write_comparator_generation_record( @@ -2024,7 +2025,12 @@ def _split_centralized_comparator_outputs( }, ] - def _parse_centralized_comparator_completion(self, text: str) -> Tuple[str, str]: + def _parse_centralized_comparator_completion( + self, + text: str, + *, + batch_item: Optional[Dict[str, Any]] = None, + ) -> Tuple[str, str]: auxiliary = self._extract_tagged_section(text, "auxiliary") main = self._extract_tagged_section(text, "main") if auxiliary is not None and main is not None: @@ -2034,10 +2040,35 @@ def _parse_centralized_comparator_completion(self, text: str) -> Tuple[str, str] ) auxiliary, main = self._extract_labeled_centralized_sections(text) + if auxiliary is not None and main is not None: + return ( + self._clean_centralized_section(auxiliary), + self._clean_centralized_section(main), + ) + + auxiliary, main = self._extract_json_centralized_sections(text) + if auxiliary is not None and main is not None: + return ( + self._clean_centralized_section(auxiliary), + self._clean_centralized_section(main), + ) + + auxiliary, main = self._extract_fenced_code_centralized_sections(text) + if auxiliary is not None and main is not None: + return ( + self._clean_centralized_section(auxiliary), + self._clean_centralized_section(main), + ) + + auxiliary, main = self._extract_function_centralized_sections( + text, + batch_item=batch_item, + ) if auxiliary is None or main is None: raise ValueError( "Centralized comparator output could not be parsed. Expected " - "... and
...
sections." + "..., Auxiliary/Main sections, JSON, " + "two code blocks, or def aux plus the task entry function." ) return ( self._clean_centralized_section(auxiliary), @@ -2075,6 +2106,114 @@ def _extract_labeled_centralized_sections( auxiliary = text[auxiliary_label.end() :] return auxiliary, main + @staticmethod + def _extract_json_centralized_sections( + text: str, + ) -> Tuple[Optional[str], Optional[str]]: + value = text.strip() + fence_match = re.fullmatch( + r"```(?:json)?\s*(.*?)```", value, flags=re.IGNORECASE | re.DOTALL + ) + if fence_match: + value = fence_match.group(1).strip() + + candidates = [value] + start = value.find("{") + end = value.rfind("}") + if start >= 0 and end > start: + candidates.append(value[start : end + 1]) + + for candidate in candidates: + try: + parsed = json.loads(candidate) + except (TypeError, json.JSONDecodeError): + continue + if not isinstance(parsed, dict): + continue + auxiliary = parsed.get("auxiliary") + if auxiliary is None: + auxiliary = parsed.get("aux") + main = parsed.get("main") + if auxiliary is not None and main is not None: + return str(auxiliary), str(main) + return None, None + + @staticmethod + def _extract_fenced_code_centralized_sections( + text: str, + ) -> Tuple[Optional[str], Optional[str]]: + blocks = re.findall( + r"```(?:python)?\s*(.*?)```", + text, + flags=re.IGNORECASE | re.DOTALL, + ) + if len(blocks) >= 2: + return blocks[0], blocks[1] + return None, None + + @classmethod + def _extract_function_centralized_sections( + cls, + text: str, + *, + batch_item: Optional[Dict[str, Any]] = None, + ) -> Tuple[Optional[str], Optional[str]]: + entry_point = None + if isinstance(batch_item, dict): + value = batch_item.get("entry_point") + if value is not None: + entry_point = str(value).strip() or None + + auxiliary = cls._extract_python_function(text, "aux") + main = cls._extract_python_function(text, entry_point) if entry_point else None + if main is None: + for name in cls._python_function_names(text): + if name != "aux": + main = cls._extract_python_function(text, name) + break + return auxiliary, main + + @staticmethod + def _python_function_names(text: str) -> List[str]: + return [ + match.group(1) + for match in re.finditer( + r"(?m)^\s*def\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(", + text, + ) + ] + + @staticmethod + def _extract_python_function( + text: str, function_name: Optional[str] + ) -> Optional[str]: + if not function_name: + return None + pattern = re.compile(rf"(?m)^([ \t]*)def\s+{re.escape(function_name)}\s*\(") + match = pattern.search(text) + if match is None: + return None + + lines = text[match.start() :].splitlines() + if not lines: + return None + base_indent = len(lines[0]) - len(lines[0].lstrip(" \t")) + selected = [lines[0]] + for line in lines[1:]: + stripped = line.strip() + indent = len(line) - len(line.lstrip(" \t")) + if stripped and indent <= base_indent: + if re.match(r"(def|class)\s+", stripped): + break + if re.match(r"?:?$", stripped, re.I): + break + if re.match( + r"(?:auxiliary|aux|main)\s*(?:code|output)?\s*:", stripped, re.I + ): + break + selected.append(line) + return "\n".join(selected).strip() + @staticmethod def _clean_centralized_section(text: str) -> str: value = text.strip() From 836e426399879bc7c26cac11b5d5d1480d64a648 Mon Sep 17 00:00:00 2001 From: N!no Date: Sat, 4 Jul 2026 13:30:17 -0400 Subject: [PATCH 3/4] Update iterative.py --- comlrl/trainers/preference/iterative.py | 125 +++++++++++++++++------- 1 file changed, 91 insertions(+), 34 deletions(-) diff --git a/comlrl/trainers/preference/iterative.py b/comlrl/trainers/preference/iterative.py index 9f1cbf2..e9eb0b4 100644 --- a/comlrl/trainers/preference/iterative.py +++ b/comlrl/trainers/preference/iterative.py @@ -1927,6 +1927,7 @@ def _build_centralized_comparator_prompt(self, batch_item: Dict[str, Any]) -> st ) auxiliary_prompt = self.formatters[0](batch_item) main_prompt = self.formatters[1](batch_item) + entry_signature = self._centralized_main_signature(batch_item) return f"""You are acting as one centralized coordinator for two code-generation agents. Your job is to produce the exact outputs that two decentralized agents would submit: @@ -1942,6 +1943,7 @@ def _build_centralized_comparator_prompt(self, batch_item: Dict[str, Any]) -> st - Return exactly two complete code snippets: one for Auxiliary and one for Main. - The Auxiliary snippet must define the helper function aux(...). - The Main snippet must define only the required main function from the problem. +- Do not copy placeholders such as required_function, ..., result, or pass. - Do not include explanations, examples, tests, markdown fences, or extra text. - Use the exact section tags below so the parser can split the two outputs. @@ -1952,12 +1954,27 @@ def aux(...):
-def required_function(...): +{entry_signature}: # your main code here return result
""" + @staticmethod + def _centralized_main_signature(batch_item: Dict[str, Any]) -> str: + entry_point = str(batch_item.get("entry_point") or "").strip() + prompt = str(batch_item.get("prompt") or "") + if entry_point: + pattern = rf"def\s+{re.escape(entry_point)}\s*\(([^)]*)\)" + match = re.search(pattern, prompt) + if match: + return f"def {entry_point}({match.group(1).strip()})" + return f"def {entry_point}(...)" + match = re.search(r"def\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(([^)]*)\)", prompt) + if match: + return f"def {match.group(1)}({match.group(2).strip()})" + return "def main(...)" + def _split_centralized_comparator_outputs( self, completions: Sequence[str], @@ -1992,9 +2009,10 @@ def _split_centralized_comparator_outputs( "parse_error": str(exc), }, ) - raise + continue auxiliary_completions.append(auxiliary) main_completions.append(main) + parse_warning = self._centralized_parse_warning(auxiliary, main) self._write_comparator_generation_record( iteration_idx, { @@ -2009,9 +2027,16 @@ def _split_centralized_comparator_outputs( "parsed_main": main, "parse_ok": True, "parse_error": None, + "parse_warning": parse_warning, }, ) + if not auxiliary_completions: + raise ValueError( + "Centralized comparator produced no parseable candidates. See " + "comparator_generations JSONL for raw completions." + ) + return [ { "prompts": [prompt], @@ -2031,49 +2056,74 @@ def _parse_centralized_comparator_completion( *, batch_item: Optional[Dict[str, Any]] = None, ) -> Tuple[str, str]: + partial_auxiliary: Optional[str] = None + partial_main: Optional[str] = None + + def _remember( + auxiliary: Optional[str], + main: Optional[str], + ) -> Optional[Tuple[str, str]]: + nonlocal partial_auxiliary, partial_main + clean_auxiliary = ( + self._clean_centralized_section(auxiliary) + if auxiliary is not None + else None + ) + clean_main = ( + self._clean_centralized_section(main) if main is not None else None + ) + if clean_auxiliary is not None and partial_auxiliary is None: + partial_auxiliary = clean_auxiliary + if clean_main is not None and partial_main is None: + partial_main = clean_main + if clean_auxiliary is not None and clean_main is not None: + return clean_auxiliary, clean_main + return None + auxiliary = self._extract_tagged_section(text, "auxiliary") main = self._extract_tagged_section(text, "main") - if auxiliary is not None and main is not None: - return ( - self._clean_centralized_section(auxiliary), - self._clean_centralized_section(main), - ) + result = _remember(auxiliary, main) + if result is not None: + return result auxiliary, main = self._extract_labeled_centralized_sections(text) - if auxiliary is not None and main is not None: - return ( - self._clean_centralized_section(auxiliary), - self._clean_centralized_section(main), - ) + result = _remember(auxiliary, main) + if result is not None: + return result auxiliary, main = self._extract_json_centralized_sections(text) - if auxiliary is not None and main is not None: - return ( - self._clean_centralized_section(auxiliary), - self._clean_centralized_section(main), - ) + result = _remember(auxiliary, main) + if result is not None: + return result auxiliary, main = self._extract_fenced_code_centralized_sections(text) - if auxiliary is not None and main is not None: - return ( - self._clean_centralized_section(auxiliary), - self._clean_centralized_section(main), - ) + result = _remember(auxiliary, main) + if result is not None: + return result auxiliary, main = self._extract_function_centralized_sections( text, batch_item=batch_item, ) - if auxiliary is None or main is None: - raise ValueError( - "Centralized comparator output could not be parsed. Expected " - "..., Auxiliary/Main sections, JSON, " - "two code blocks, or def aux plus the task entry function." - ) - return ( - self._clean_centralized_section(auxiliary), - self._clean_centralized_section(main), - ) + result = _remember(auxiliary, main) + if result is not None: + return result + + if partial_auxiliary is not None or partial_main is not None: + return partial_auxiliary or "", partial_main or "" + + return "", "" + + @staticmethod + def _centralized_parse_warning(auxiliary: str, main: str) -> Optional[str]: + missing = [] + if not auxiliary.strip(): + missing.append("auxiliary") + if not main.strip(): + missing.append("main") + if not missing: + return None + return "missing_" + "_and_".join(missing) @staticmethod def _extract_tagged_section(text: str, tag: str) -> Optional[str]: @@ -2096,6 +2146,10 @@ def _extract_labeled_centralized_sections( text, ) if auxiliary_label is None or main_label is None: + if auxiliary_label is not None: + return text[auxiliary_label.end() :], None + if main_label is not None: + return None, text[main_label.end() :] return None, None if auxiliary_label.start() < main_label.start(): @@ -2134,8 +2188,11 @@ def _extract_json_centralized_sections( if auxiliary is None: auxiliary = parsed.get("aux") main = parsed.get("main") - if auxiliary is not None and main is not None: - return str(auxiliary), str(main) + if auxiliary is not None or main is not None: + return ( + str(auxiliary) if auxiliary is not None else None, + str(main) if main is not None else None, + ) return None, None @staticmethod From 5b133e82dde4177d925f8159f1ced796338825af Mon Sep 17 00:00:00 2001 From: N!no Date: Sat, 4 Jul 2026 15:46:26 -0400 Subject: [PATCH 4/4] rm log of raw comparator api input --- comlrl/trainers/preference/iterative.py | 281 ++++++++---------------- 1 file changed, 86 insertions(+), 195 deletions(-) diff --git a/comlrl/trainers/preference/iterative.py b/comlrl/trainers/preference/iterative.py index e9eb0b4..c5bd7d1 100644 --- a/comlrl/trainers/preference/iterative.py +++ b/comlrl/trainers/preference/iterative.py @@ -259,7 +259,6 @@ class MADPOIterConfig(MADPOConfig): preference_replay_sample_size: Optional[int] = None preference_replay_dir: Optional[str] = None log_reward_distribution: bool = False - log_comparator_generations: bool = False policy_checkpoint_dir: Optional[str] = None comparator_policy: str = "current" comparator_generation_mode: str = "decentralized" @@ -306,7 +305,6 @@ class MARLHFIterConfig(MARLHFConfig): preference_replay_sample_size: Optional[int] = None preference_replay_dir: Optional[str] = None log_reward_distribution: bool = False - log_comparator_generations: bool = False preference_scoring_reward: str = "task" policy_checkpoint_dir: Optional[str] = None comparator_policy: str = "current" @@ -1187,62 +1185,6 @@ def _reward_distribution_dir(self) -> str: self._reward_distribution_dir_path = path return path - def _comparator_generation_dir(self) -> str: - cached = getattr(self, "_comparator_generation_dir_path", None) - if cached: - return cached - - path = None - if isinstance(self.wandb_config, dict): - output_dir = self.wandb_config.get("output_dir") - if output_dir: - path = os.path.join(str(output_dir), "comparator_generations") - else: - sections = self.wandb_config.get("config_sections") or {} - output_section = ( - sections.get("output") if isinstance(sections, dict) else {} - ) - base_dir = None - if isinstance(output_section, dict): - base_dir = output_section.get("base_dir") - base_dir = base_dir or self.wandb_config.get("dir") - if base_dir: - job_id = os.environ.get("SLURM_JOB_ID") - path = ( - os.path.join( - str(base_dir), f"job_{job_id}", "comparator_generations" - ) - if job_id - else os.path.join(str(base_dir), "comparator_generations") - ) - if not path: - path = os.path.join(os.getcwd(), "comparator_generations") - - path = os.path.abspath(str(path)) - os.makedirs(path, exist_ok=True) - self._comparator_generation_dir_path = path - return path - - def _write_comparator_generation_record( - self, - iteration_idx: int, - record: Dict[str, Any], - ) -> None: - if not getattr(self.args, "log_comparator_generations", False): - return - path = os.path.join( - self._comparator_generation_dir(), - f"iteration_{iteration_idx + 1:04d}.jsonl", - ) - payload = { - "iteration": int(iteration_idx + 1), - "comparator_policy": self.args.comparator_policy, - "comparator_generation_mode": self.args.comparator_generation_mode, - **record, - } - with open(path, "a", encoding="utf-8") as handle: - handle.write(json.dumps(self._jsonable(payload), ensure_ascii=False) + "\n") - def _write_iteration_preference_pairs( self, iteration_idx: int, @@ -1763,54 +1705,30 @@ def _generate_comparator_outputs_for_item( ) if self.args.comparator_policy == "api": - outputs = self._generate_api_outputs_for_item( + return self._generate_api_outputs_for_item( batch_item, num_candidates=num_candidates, ) - self._log_decentralized_comparator_outputs( - outputs, - batch_item=batch_item, - iteration_idx=iteration_idx, - ) - return outputs if self.args.comparator_policy == "current": - outputs = self._generate_policy_outputs_for_item( + return self._generate_policy_outputs_for_item( self.agents, batch_item, num_candidates=num_candidates, **kwargs, ) - self._log_decentralized_comparator_outputs( - outputs, - batch_item=batch_item, - iteration_idx=iteration_idx, - ) - return outputs if self.args.comparator_policy == "history": - outputs = self._generate_history_policy_outputs_for_item( + return self._generate_history_policy_outputs_for_item( batch_item, iteration_idx=iteration_idx, num_candidates=num_candidates, **kwargs, ) - self._log_decentralized_comparator_outputs( - outputs, - batch_item=batch_item, - iteration_idx=iteration_idx, - ) - return outputs - outputs = self._generate_policy_outputs_for_item( + return self._generate_policy_outputs_for_item( self._get_comparator_agents(), batch_item, num_candidates=num_candidates, **kwargs, ) - self._log_decentralized_comparator_outputs( - outputs, - batch_item=batch_item, - iteration_idx=iteration_idx, - ) - return outputs def _generate_policy_outputs_for_item( self, @@ -1852,56 +1770,55 @@ def _generate_centralized_comparator_outputs_for_item( completions, batch_item=batch_item, prompt=prompt, - iteration_idx=iteration_idx, ) if self.args.comparator_policy == "current": - return self._generate_centralized_policy_outputs_for_item( - self.agents, + return self._generate_centralized_policy_output_for_agent( + self.agents[self._centralized_comparator_agent_index()], batch_item, prompt=prompt, - iteration_idx=iteration_idx, num_candidates=num_candidates, **kwargs, ) if self.args.comparator_policy == "history": checkpoint_dir = self._history_policy_checkpoint_path(iteration_idx) - comparator_agents = self._load_policy_checkpoint_agents(checkpoint_dir) + agent_idx = self._centralized_comparator_agent_index() + comparator_agent = self._load_single_policy_checkpoint_agent( + checkpoint_dir, + agent_idx, + ) try: - return self._generate_centralized_policy_outputs_for_item( - comparator_agents, + return self._generate_centralized_policy_output_for_agent( + comparator_agent, batch_item, prompt=prompt, - iteration_idx=iteration_idx, num_candidates=num_candidates, **kwargs, ) finally: - del comparator_agents + del comparator_agent gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() - return self._generate_centralized_policy_outputs_for_item( - self._get_comparator_agents(), + return self._generate_centralized_policy_output_for_agent( + self._get_centralized_comparator_agent(), batch_item, prompt=prompt, - iteration_idx=iteration_idx, num_candidates=num_candidates, **kwargs, ) - def _generate_centralized_policy_outputs_for_item( + def _generate_centralized_policy_output_for_agent( self, - policy_agents: Sequence[Any], + policy_agent: Any, batch_item: Dict[str, Any], *, prompt: str, - iteration_idx: int, num_candidates: int, **kwargs, ) -> List[Dict[str, Any]]: agent_idx = self._centralized_comparator_agent_index() generation_output = self._generate_completions( - policy_agents[agent_idx], + policy_agent, [batch_item], agent_idx=agent_idx, num_return_sequences=num_candidates, @@ -1914,7 +1831,6 @@ def _generate_centralized_policy_outputs_for_item( completions, batch_item=batch_item, prompt=prompt, - iteration_idx=iteration_idx, ) def _centralized_comparator_agent_index(self) -> int: @@ -1981,61 +1897,20 @@ def _split_centralized_comparator_outputs( *, batch_item: Dict[str, Any], prompt: str, - iteration_idx: int, ) -> List[Dict[str, Any]]: auxiliary_completions: List[str] = [] main_completions: List[str] = [] - for candidate_idx, completion in enumerate(completions): + for completion in completions: raw_completion = str(completion) - try: - auxiliary, main = self._parse_centralized_comparator_completion( - raw_completion, - batch_item=batch_item, - ) - except Exception as exc: - self._write_comparator_generation_record( - iteration_idx, - { - "candidate_index": int(candidate_idx), - "centralized_agent_index": ( - self._centralized_comparator_agent_index() - ), - "batch_item": self._jsonable(batch_item), - "prompt": prompt, - "raw_completion": raw_completion, - "parsed_auxiliary": None, - "parsed_main": None, - "parse_ok": False, - "parse_error": str(exc), - }, - ) - continue + auxiliary, main = self._parse_centralized_comparator_completion( + raw_completion, + batch_item=batch_item, + ) auxiliary_completions.append(auxiliary) main_completions.append(main) - parse_warning = self._centralized_parse_warning(auxiliary, main) - self._write_comparator_generation_record( - iteration_idx, - { - "candidate_index": int(candidate_idx), - "centralized_agent_index": ( - self._centralized_comparator_agent_index() - ), - "batch_item": self._jsonable(batch_item), - "prompt": prompt, - "raw_completion": raw_completion, - "parsed_auxiliary": auxiliary, - "parsed_main": main, - "parse_ok": True, - "parse_error": None, - "parse_warning": parse_warning, - }, - ) if not auxiliary_completions: - raise ValueError( - "Centralized comparator produced no parseable candidates. See " - "comparator_generations JSONL for raw completions." - ) + raise ValueError("Centralized comparator produced no candidates.") return [ { @@ -2114,17 +1989,6 @@ def _remember( return "", "" - @staticmethod - def _centralized_parse_warning(auxiliary: str, main: str) -> Optional[str]: - missing = [] - if not auxiliary.strip(): - missing.append("auxiliary") - if not main.strip(): - missing.append("main") - if not missing: - return None - return "missing_" + "_and_".join(missing) - @staticmethod def _extract_tagged_section(text: str, tag: str) -> Optional[str]: pattern = rf"<\s*{tag}\s*>(.*?)<\s*/\s*{tag}\s*>" @@ -2281,40 +2145,6 @@ def _clean_centralized_section(text: str) -> str: value = fence_match.group(1).strip() return value - def _log_decentralized_comparator_outputs( - self, - outputs: Sequence[Dict[str, Any]], - *, - batch_item: Dict[str, Any], - iteration_idx: int, - ) -> None: - if not getattr(self.args, "log_comparator_generations", False): - return - for agent_idx, output in enumerate(outputs): - prompts = output.get("prompts") if isinstance(output, dict) else None - prompt = prompts[0] if isinstance(prompts, list) and prompts else None - completions = ( - output.get("completions") if isinstance(output, dict) else None - ) - candidate_completions = ( - completions[0] if isinstance(completions, list) and completions else [] - ) - for candidate_idx, completion in enumerate(candidate_completions): - self._write_comparator_generation_record( - iteration_idx, - { - "agent_idx": int(agent_idx), - "candidate_index": int(candidate_idx), - "batch_item": self._jsonable(batch_item), - "prompt": prompt, - "raw_completion": str(completion), - "parsed_auxiliary": None, - "parsed_main": None, - "parse_ok": True, - "parse_error": None, - }, - ) - def _generate_history_policy_outputs_for_item( self, batch_item: Dict[str, Any], @@ -2703,6 +2533,39 @@ def _get_comparator_agents(self) -> Sequence[Any]: self._comparator_agents = self._load_comparator_agents() return self._comparator_agents + def _get_centralized_comparator_agent(self) -> Any: + if self.args.comparator_policy == "current": + return self.agents[self._centralized_comparator_agent_index()] + if self.args.comparator_policy == "history": + raise ValueError("History comparator agents must be loaded per iteration.") + if getattr(self, "_centralized_comparator_agent", None) is None: + source = self._centralized_comparator_source() + self._centralized_comparator_agent = self._load_single_frozen_policy_agent( + source, + self._centralized_comparator_agent_index(), + ) + return self._centralized_comparator_agent + + def _centralized_comparator_source(self) -> Any: + if self.args.comparator_agents is not None: + sources = self.args.comparator_agents + if isinstance(sources, (str, bytes)) or not isinstance(sources, Sequence): + raise ValueError("comparator_agents must be a non-empty sequence.") + if len(sources) == 1: + return list(sources)[0] + if len(sources) == self.num_agents: + return list(sources)[self._centralized_comparator_agent_index()] + raise ValueError( + "comparator_agents length must be 1 or num_agents when " + "comparator_generation_mode='centralized'." + ) + if self.args.comparator_model_name is None: + raise ValueError( + "comparator_model_name or comparator_agents is required when " + "comparator_policy='model'." + ) + return self.args.comparator_model_name + def _load_comparator_agents(self) -> List[Any]: comparator_sources, _ = resolve_model_sources( kind="comparator_agents", @@ -2727,6 +2590,19 @@ def _load_policy_checkpoint_agents(self, checkpoint_dir: str) -> List[Any]: ) return self._load_frozen_policy_agents(checkpoint_sources) + def _load_single_policy_checkpoint_agent( + self, + checkpoint_dir: str, + agent_idx: int, + ) -> Any: + checkpoint_source = os.path.join(checkpoint_dir, f"agent_{agent_idx}") + if not os.path.isdir(checkpoint_source): + raise ValueError( + "Missing comparator history checkpoint agent directory: " + f"{checkpoint_source}." + ) + return self._load_single_frozen_policy_agent(checkpoint_source, agent_idx) + def _load_frozen_policy_agents(self, sources: Sequence[Any]) -> List[Any]: comparator_devices = self._resolve_comparator_devices() model_kwargs = self._comparator_model_kwargs() @@ -2747,6 +2623,21 @@ def _load_frozen_policy_agents(self, sources: Sequence[Any]) -> List[Any]: return comparator_agents + def _load_single_frozen_policy_agent(self, source: Any, agent_idx: int) -> Any: + model_kwargs = self._comparator_model_kwargs() + comparator_agent = ( + AutoModelForCausalLM.from_pretrained(source, **model_kwargs) + if isinstance(source, str) + else source + ) + comparator_device = self._resolve_comparator_devices()[agent_idx] + comparator_agent.to(comparator_device) + comparator_agent.eval() + for param in comparator_agent.parameters(): + param.requires_grad = False + apply_tokenizer_specials(self.tokenizers[agent_idx], [comparator_agent]) + return comparator_agent + def _resolve_comparator_devices(self) -> List[torch.device]: return DeviceScheduler.resolve_devices( self.args.comparator_devices or getattr(self.args, "agent_devices", None),