Merge branch 'agentjet_codebase_backup' into HEAD

Author: binary-husky

.gitignore

@@ -191,3 +191,9 @@ cocktail_results_*
 PRINCIPLE.md
 TODO.md
 plot_mean_reward.py
+appworld_swarm_results_token
+appworld_swarm_results_text
+appworld_swarm_results
+TODO
+.trash.*
+tutorial/opencode_build_aime/auto_research/results/*

ajet/backbone/main_verl.py

@@ -59,9 +59,7 @@ def run_ppo(config: DictConfig) -> None:
     if not ray.is_initialized():
         # this is for local ray cluster
         runtime_env = get_runtime_env(config)
-        ray.init(
-            runtime_env=runtime_env,
-        )
+        ray.init(runtime_env=runtime_env)
 
     def on_shutdown():
         if ray.is_initialized():

ajet/backbone/trainer_verl.py

@@ -131,8 +131,6 @@ def union_gen_batch_via_task_id(tasks, batch: DataProto, gen_batch_output: DataP
                 task_id_counter[tid] += 1
             else:
                 task_id_counter[tid] = 1
-            current_id = task_id_counter[tid]
-            gen_batch_output.non_tensor_batch['rollout_ids'][i] = f"T{tid}R{current_id}"
         logger.info(f'task_id_counter: {task_id_counter}')
         return gen_batch_output
 
@@ -166,6 +164,117 @@ def import_or_export_data_proto(batch: DataProto, direction: str = "export", fil
     else:
         raise ValueError(f"direction must be 'import' or 'export', got '{direction}'")
 
+def compute_grpo_episode_level_outcome_advantage(
+    token_level_rewards: torch.Tensor,
+    response_mask: torch.Tensor,
+    index: np.ndarray,
+    episode_index: np.ndarray,
+    norm_adv_by_std_in_grpo: bool = True,
+    epsilon: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """GRPO outcome advantage with the baseline computed at *episode* scope.
+
+    Mirrors ``verl.trainer.ppo.core_algos.compute_grpo_outcome_advantage`` but,
+    instead of treating every sample equally when forming the per-task (``uid``)
+    baseline, it first reduces every episode (``episode_uuids``) to its mean
+    scalar reward and then computes the task baseline mean/std over those
+    per-episode means. This way an episode that produced many samples does not
+    dominate the baseline of an episode that produced few.
+
+    Example (matches the documented behaviour):
+        task T -> episode 1 (2 samples, reward 1) + episode 2 (1 sample, reward 0)
+        sample scope baseline  = (1 + 1 + 0) / 3 = 0.667
+        episode scope baseline = (mean[1, 1] + mean[0]) / 2 = (1 + 0) / 2 = 0.5
+
+    Args:
+        token_level_rewards: (bsz, response_length) reward tensor.
+        response_mask: (bsz, response_length) mask of trainable response tokens.
+        index: per-sample task id (``non_tensor_batch["uid"]``).
+        episode_index: per-sample episode id (``non_tensor_batch["episode_uuids"]``).
+        norm_adv_by_std_in_grpo: divide the centred reward by the (episode-level)
+            group std when True, otherwise only subtract the group mean.
+        epsilon: numerical-stability term added to the std denominator.
+
+    Returns:
+        (advantages, returns) - both (bsz, response_length); identical, as in GRPO.
+    """
+    scores = (token_level_rewards * response_mask).sum(dim=-1)  # (bsz,) scalar reward
+    bsz = scores.shape[0]
+
+    # 1) reduce each episode to its mean scalar reward
+    episode_score_sum: dict = defaultdict(float)
+    episode_score_cnt: dict = defaultdict(int)
+    for i in range(bsz):
+        ep = episode_index[i]
+        episode_score_sum[ep] += scores[i].item()
+        episode_score_cnt[ep] += 1
+    episode_mean = {ep: episode_score_sum[ep] / episode_score_cnt[ep] for ep in episode_score_sum}
+
+    # 2) collect, per task, the set of distinct episodes it produced
+    task2episodes: dict = defaultdict(dict)  # use dict as ordered set
+    for i in range(bsz):
+        task2episodes[index[i]][episode_index[i]] = None
+
+    # 3) per-task baseline = mean/std over the per-episode means.
+    #    Single-episode tasks are degenerate -> follow verl's convention
+    #    (mean=0, std=1) so the advantage reduces to the raw score.
+    task_mean: dict = {}
+    task_std: dict = {}
+    for task, episodes in task2episodes.items():
+        vals = torch.tensor([episode_mean[ep] for ep in episodes], dtype=torch.float32)
+        if vals.numel() == 1:
+            task_mean[task] = torch.tensor(0.0)
+            task_std[task] = torch.tensor(1.0)
+        else:
+            task_mean[task] = vals.mean()
+            task_std[task] = vals.std()
+
+    # 4) centre (and optionally normalise) every sample against its task baseline
+    adv = scores.clone()
+    for i in range(bsz):
+        task = index[i]
+        if norm_adv_by_std_in_grpo:
+            adv[i] = (scores[i] - task_mean[task]) / (task_std[task] + epsilon)
+        else:
+            adv[i] = scores[i] - task_mean[task]
+
+    adv = adv.unsqueeze(-1) * response_mask
+    return adv, adv
+
+
+def compute_episode_level_loss_weight(data: DataProto) -> torch.Tensor:
+    """Per-token loss weight that makes every episode contribute equally.
+
+    Each sample belonging to an episode (same ``non_tensor_batch["episode_uuids"]``)
+    that produced ``N`` samples receives weight ``1 / N``. The weights of all
+    samples of one episode therefore sum to 1, so an episode that emitted many
+    samples does not contribute more to the loss than one that emitted few.
+
+    The weight is broadcast across the response dimension so it has the **same
+    shape as ``advantages``** ((bsz, response_length)); this lets it multiply
+    both the per-token policy-gradient term and the per-token KL term directly.
+
+    Returns:
+        A (bsz, response_length) tensor (matching ``data.batch["advantages"]``
+        dtype/device) of per-token loss weights, constant along the response
+        dimension for a given sample.
+    """
+    episode_index = data.non_tensor_batch["episode_uuids"]
+    bsz = len(episode_index)
+    episode_count: dict = defaultdict(int)
+    for ep in episode_index:
+        episode_count[ep] += 1
+    advantages = data.batch["advantages"]  # (bsz, response_length)
+    per_sample = torch.tensor(
+        [1.0 / episode_count[episode_index[i]] for i in range(bsz)],
+        dtype=advantages.dtype,
+        device=advantages.device,
+    )
+    # broadcast per-sample weight to the same shape as advantages
+    weights = per_sample.view(-1, 1) * torch.ones_like(advantages)
+    return weights
+
+
 def compute_advantage(
     data: DataProto,
     adv_estimator: AdvantageEstimator,
@@ -174,6 +283,7 @@ def compute_advantage(
     num_repeat: int = 1,
     norm_adv_by_std_in_grpo: bool = True,
     config: Optional[AlgoConfig] = None,
+    advantage_estimation_episode_level: bool = False,
 ) -> DataProto:
     """Compute advantage estimates for policy optimization.
 
@@ -189,13 +299,21 @@ def compute_advantage(
         norm_adv_by_std_in_grpo (bool, optional): Whether to normalize advantages by standard deviation in
             GRPO. Defaults to True.
         config (dict, optional): Configuration dictionary for algorithm settings. Defaults to None.
+        advantage_estimation_episode_level (bool, optional): When True (and using the GRPO estimator),
+            the GRPO baseline is computed at episode scope instead of sample scope so every episode
+            contributes equally regardless of how many samples it produced. Defaults to False.
 
     Returns:
         DataProto: The updated data with computed advantages and returns.
     """
     # Back-compatible with trainers that do not compute response mask in fit
     if "response_mask" not in data.batch.keys():
         data.batch["response_mask"] = compute_response_mask(data)
+    if advantage_estimation_episode_level and adv_estimator != AdvantageEstimator.GRPO:
+        raise NotImplementedError(
+            "ajet.trainer_common.advantage_estimation_episode_level is only "
+            f"supported with the GRPO advantage estimator, got {adv_estimator}."
+        )
     # prepare response group
     if adv_estimator == AdvantageEstimator.GAE:
         # Compute advantages and returns using Generalized Advantage Estimation (GAE)
@@ -222,13 +340,30 @@ def compute_advantage(
         response_length = grpo_calculation_mask.size(1)
         # This mask is the one intended for GRPO
         grpo_calculation_mask = data.batch["loss_mask"][:, -response_length:]
-        # Call compute_grpo_outcome_advantage with parameters matching its definition
-        advantages, returns = core_algos.compute_grpo_outcome_advantage(
-            token_level_rewards=data.batch["token_level_rewards"],
-            response_mask=grpo_calculation_mask,
-            index=data.non_tensor_batch["uid"],
-            norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
-        )
+        if advantage_estimation_episode_level:
+            # Episode-scope baseline: every episode contributes equally to the
+            # per-task baseline regardless of how many samples it produced.
+            if "episode_uuids" not in data.non_tensor_batch:
+                raise KeyError(
+                    "advantage_estimation_episode_level is enabled but "
+                    "non_tensor_batch['episode_uuids'] is missing; cannot identify "
+                    "same-episode samples."
+                )
+            advantages, returns = compute_grpo_episode_level_outcome_advantage(
+                token_level_rewards=data.batch["token_level_rewards"],
+                response_mask=grpo_calculation_mask,
+                index=data.non_tensor_batch["uid"],
+                episode_index=data.non_tensor_batch["episode_uuids"],
+                norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+            )
+        else:
+            # Call compute_grpo_outcome_advantage with parameters matching its definition
+            advantages, returns = core_algos.compute_grpo_outcome_advantage(
+                token_level_rewards=data.batch["token_level_rewards"],
+                response_mask=grpo_calculation_mask,
+                index=data.non_tensor_batch["uid"],
+                norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+            )
         data.batch["advantages"] = advantages
         data.batch["returns"] = returns
     else:
@@ -692,6 +827,13 @@ def fit(self):  # noqa: C901
                             "norm_adv_by_std_in_grpo", True
                         )  # GRPO adv normalization factor
 
+                        # [AJET] episode-scope advantage baseline (disabled by default)
+                        advantage_estimation_episode_level = bool(
+                            self.config.ajet.trainer_common.get(
+                                "advantage_estimation_episode_level", False
+                            )
+                        )
+
                         batch = compute_advantage(
                             batch,
                             adv_estimator=self.config.algorithm.adv_estimator,
@@ -700,8 +842,26 @@ def fit(self):  # noqa: C901
                             num_repeat=self.config.ajet.rollout.num_repeat,
                             norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
                             config=self.config.algorithm,
+                            advantage_estimation_episode_level=advantage_estimation_episode_level,
                         )
 
+                        # [AJET] per-sample loss weight that makes every episode
+                        # contribute equally to the policy-gradient update
+                        # (disabled by default). Consumed in
+                        # AjetDataParallelPPOActor.update_policy.
+                        if bool(
+                            self.config.ajet.trainer_common.get(
+                                "loss_weight_normalization_episode_level", False
+                            )
+                        ):
+                            if "episode_uuids" not in batch.non_tensor_batch:
+                                raise KeyError(
+                                    "loss_weight_normalization_episode_level is enabled but "
+                                    "non_tensor_batch['episode_uuids'] is missing; cannot "
+                                    "identify same-episode samples."
+                                )
+                            batch.batch["loss_weight"] = compute_episode_level_loss_weight(batch)
+
                     # update critic
                     if self.use_critic:
                         with marked_timer("update_critic", timing_raw, color="pink"):

ajet/backbone/verl/dp_actor.py

@@ -161,6 +161,11 @@ def update_policy(self, data: DataProto):
         # Include rollout_log_probs for computing rollout_corr metrics in bypass mode
         if "rollout_log_probs" in data.batch.keys():
             select_keys.append("rollout_log_probs")
+        # [AJET] per-sample loss weight (episode-level loss normalization).
+        # Present only when ajet.trainer_common.loss_weight_normalization_episode_level
+        # is enabled; absent => every sample weighted equally (default behaviour).
+        if "loss_weight" in data.batch.keys():
+            select_keys.append("loss_weight")
 
         has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch.keys()
         non_tensor_select_keys = []
@@ -208,6 +213,20 @@ def update_policy(self, data: DataProto):
                     response_mask = model_inputs["response_mask"]
                     old_log_prob = model_inputs["old_log_probs"]
                     advantages = model_inputs["advantages"]
+                    # [AJET] Episode-level loss-weight normalization.
+                    # When ajet.trainer_common.loss_weight_normalization_episode_level
+                    # is enabled, every sample carries a per-token weight (1/N for
+                    # an episode that produced N samples), same shape as advantages.
+                    # Scaling the advantages by this positive weight scales each
+                    # sample's policy-gradient contribution by the same factor (the
+                    # clip/ratio behaviour is unchanged since the weight is a
+                    # positive per-sample constant); the same weight is applied to
+                    # the per-token KL term below, so every episode contributes
+                    # equally to the total loss.
+                    loss_weight = model_inputs.get("loss_weight", None)
+                    if loss_weight is not None:
+                        loss_weight = loss_weight.to(advantages.dtype)
+                        advantages = advantages * loss_weight
                     # [AJET] Debug logging for tensor shapes
                     input_ids = model_inputs["input_ids"]
                     _shape_msg = f'[Update Policy] -> Micro batch shape, input_ids {input_ids.shape}, response {response_mask.shape} @{micro_batch_idx}/{num_micro_batches}'
@@ -294,6 +313,12 @@ def update_policy(self, data: DataProto):
                         kld = kl_penalty(
                             logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type
                         )
+                        # [AJET] apply the per-token episode-level loss weight to
+                        # the KL term as well (same weight/shape used for the
+                        # policy-gradient term above), so each episode contributes
+                        # equally to the KL loss too.
+                        if loss_weight is not None:
+                            kld = kld * loss_weight
                         kl_loss = agg_loss(loss_mat=kld, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
 
                         policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef

ajet/backbone/warm_up.py

@@ -6,6 +6,7 @@
 import asyncio
 import logging
 import os
+from datetime import datetime
 from ajet.utils.async_utils import (
     apply_httpx_aclose_patch,
     silence_hermes_tool_parser_loggers,
@@ -14,21 +15,19 @@
 apply_httpx_aclose_patch()
 suppress_httpx_aclose_exception()
 
-def init_parallel_rollout_logger(experiment_name, experiment_dir):
+def init_parallel_rollout_logger(experiment_dir):
     """Initialize the logger with the given configuration."""
     if "PROCESS_LEVEL_WARMUP_INIT_LOGGER" in os.environ:
         return
-    os.environ["PROCESS_LEVEL_WARMUP_INIT_LOGGER"] = "1"
 
-    from datetime import datetime
+    os.environ["PROCESS_LEVEL_WARMUP_INIT_LOGGER"] = "1"
 
     from beast_logger import register_logger
 
     final_log_path = os.path.join(
         experiment_dir,
         datetime.now().strftime("%Y_%m_%d_%H_%M"),
-        # machine host name
-        os.uname().nodename,
+        os.uname().nodename, # machine host name
     )
     os.environ["BEST_LOGGER_PATH"] = final_log_path
     non_console_mods = ["rollout", "token_clip", "bad_case"]
@@ -102,6 +101,6 @@ def warm_up_process(config):
     os.environ["PROCESS_LEVEL_WARMUP_INIT"] = "1"
     experiment_name = config.ajet.experiment_name
     experiment_dir = config.ajet.experiment_dir
-    init_parallel_rollout_logger(experiment_name, experiment_dir)
+    init_parallel_rollout_logger(experiment_dir)
     warm_up_task_judge_when_needed(config)
     clean_up_tmp_ajet_dir(config)

ajet/context_tracker/multiagent_tracking.py

@@ -440,13 +440,11 @@ def patch_prompt_tokens(
 
         # HISTORY context length vs CURRENT prompt length
         if len(previous_ext_context) != len(prompt_text_split):
-            logger.bind(exception=True).error(
-                f"Length mismatch when patching prompt tokens. Previous ext context length: {len(previous_ext_context)}, prompt text split length: {len(prompt_text_split)}. Replacing all tokens."
-            )
+            logger.bind(exception=True).error(f"Length mismatch when patching prompt tokens. Previous ext context length: {len(previous_ext_context)}, prompt text split length: {len(prompt_text_split)}. Replacing all tokens.")
 
         # try to recover tokens
         if self.config.ajet.context_tracker.fix_retokenization_drift:
-            self.ensure_retokenization_perfect_match(
+            previous_ext_context = self.ensure_retokenization_perfect_match(
                 previous_ext_context,   # HISTORY
                 split_prompt_token_ids, # CURRENT
                 prompt_text_split,      # CURRENT
@@ -481,18 +479,17 @@ def ensure_retokenization_perfect_match(self, previous_ext_context, split_prompt
                 # otherwise, we throw a warning (do not worry, this causes almost no influence in the training)
                 print_dict(
                     {
-                        "expected_prompt_text": prompt_text_split[j],       # from llm_output["prompt_text"]
-                        "current_prompt_text": current_prompt_text[j],      # history prompt text converted from token_arr to text using tokenizer
-                        "expected_token_ids": vllm_token_array,         # from llm_output["prompt_token_ids"]
-                        "current_token_ids": tracker_token_array,       # from previous_ext_context[j].token_arr
+                        "expected_prompt_text": prompt_text_split[j],       # from llm_output["prompt_text"], converted directly from messages using apply_chat_template, passway (messages->apply_chat_template->text)
+                        "current_prompt_text": current_prompt_text[j],      # history prompt text converted from token_arr to text using tokenizer, passway (messages->extended_message->incremental apply_chat_template->token_arr->text)
+                        "expected_token_ids": vllm_token_array,             # from llm_output["prompt_token_ids"], passway (messages->apply_chat_template->token)
+                        "current_token_ids": tracker_token_array,           # from previous_ext_context[j].token_arr, passway (messages->extended_message->incremental apply_chat_template->token_arr)
                     },
                     mod="exception",
-                    header="Prompt token ids mismatch.",
+                    header="Prompt token ids mismatch (fixing drift by `token_arr=vllm_token_array`).",
                 )
-                # # fix drift
-                # previous_ext_context[j].token_arr = self.tokenizer(
-                #     prompt_text_split[j], return_tensors="pt", padding=False
-                # )["input_ids"].tolist()
+                previous_ext_context[j].token_arr = vllm_token_array
+        return previous_ext_context
+
 
     def process_reward(self, reward_structure: Reward):
         self.reward_structure = reward_structure