revise commit

Author: binary-husky

ajet/backbone/trainer_verl.py

@@ -17,8 +17,6 @@
 from collections import defaultdict
 from pprint import pprint
 from typing import List, Optional
-from ajet.utils.metric_helper.reward_metric_helper import compute_reward_metrics_from_cmts
-from loguru import logger as loguru_logger
 
 import hydra
 import numpy as np
@@ -56,14 +54,7 @@
 from ajet.schema.task import Task
 from ajet.task_reader import dict_to_ajet_task
 from ajet.task_rollout.native_parallel_worker import VerlRolloutManager
-from ajet.utils.save_trajectory import save_train_trajectory, save_eval_trajectory
-from ajet.utils.msg_converter import (
-    convert_grouped_steps_to_openai_format,
-    convert_ext_msg_to_openai_format,
-    agentscope_to_openai,
-    openai_to_agentscope,
-)
-from ajet.utils.metric_helper.tool_metric_helper import compute_tool_metrics_from_cmts
+from ajet.utils.metric_helper import save_trajectory_as_json_file, update_metrics
 
 def parse_reward_from_dataproto(data: DataProto, return_dict=False) -> dict | torch.Tensor:
     """
@@ -586,9 +577,6 @@ def fit(self):  # noqa: C901
                             tasks, mode="sample", epoch=f"train.{epoch}"
                         )
                         logger.info("=" * 10 + "end fit rollout" + "=" * 10)
-                        
-                        if self.config.ajet.trainer_common.save_trajectory:
-                            save_train_trajectory(context_tracker_arr, self.global_steps)
                         logger.info("begin to convert context_tracker_arr to dataproto")
                         gen_batch_output = self.parallel_env.to_dataproto(context_tracker_arr)
                         logger.info("end convertion")
@@ -614,14 +602,8 @@ def fit(self):  # noqa: C901
                                 ),
                             }
                         )
-                        from ajet.utils.metric_helper.tool_metric_helper import compute_tool_metrics_from_trajectories
-                        from ajet.utils.metric_helper.reward_metric_helper import compute_reward_metrics_from_trajectories
-                        tool_metrics = compute_tool_metrics_from_trajectories(context_tracker_arr)
-                        reward_metrics = compute_reward_metrics_from_trajectories(context_tracker_arr)
-                        if tool_metrics:
-                            metrics.update(tool_metrics)
-                        if reward_metrics:
-                            metrics.update(reward_metrics)
+                        save_trajectory_as_json_file(context_tracker_arr, self.global_steps, self.config, prefix="train")
+                        update_metrics(context_tracker_arr, metrics)
                         if self.config.ajet.execute_test:  # apply a test probe
                             from swanlab.data.run.main import get_run
 
@@ -1049,10 +1031,6 @@ def eval_dataset(self, target_dataset, target_dataset_name, mode, epoch):
         for ctx_tracker in ctx_trackers:
             ctx_tracker.generate_log()
 
-        # save eval trajectories
-        if self.config.ajet.trainer_common.save_trajectory:
-            save_eval_trajectory(ctx_trackers, self.global_steps)
-
         rewards = [ctx_tracker.reward_structure.raw_reward for ctx_tracker in ctx_trackers]
         num_tasks = len(task_results)
         assert num_tasks == len(ctx_trackers) // pass_n
@@ -1068,12 +1046,8 @@ def eval_dataset(self, target_dataset, target_dataset_name, mode, epoch):
             f"TGC@{pass_n}-all-pass": num_all_success_tasks / num_tasks,
             "mean_reward": sum(rewards) / len(rewards) if rewards else 0,
         }
-        reward_metrics = compute_reward_metrics_from_cmts(ctx_trackers, print_debug=True)  
-        tool_metrics = compute_tool_metrics_from_cmts(ctx_trackers) 
-        if tool_metrics:
-            val_metrics.update(reward_metrics)
-        if reward_metrics:
-            val_metrics.update(tool_metrics)
+        save_trajectory_as_json_file(ctx_trackers, self.global_steps, self.config, prefix="eval")
+        update_metrics(ctx_trackers, val_metrics)
         print_dict(
             val_metrics,
             narrow=True,

ajet/context_tracker/base_tracker.py

@@ -1,5 +1,6 @@
 from typing import List, Tuple, Union
 from typing import List, Union, Tuple, Dict, Optional, Any
+from ajet.schema.task import WorkflowTask
 
 from ajet.schema.extended_msg import (
     INVALID_LOG_PROB_VALUE,
@@ -111,10 +112,14 @@ def replace_token_ids(
 
 
 class BaseTracker(object):
-    def __init__(self, config, tokenizer, **kwargs):
-        self.task_batch_index = kwargs.get("task_batch_index", "undefined")
-        self.task_tag = kwargs.get("task_tag", "undefined")
-        self.task_id = kwargs.get("task_id", "undefined")
+    def __init__(self, config, tokenizer, workflow_task: WorkflowTask, **kwargs):
+
+        self.workflow_task = workflow_task
+        self.task_batch_index = self.workflow_task.task_batch_index
+        self.task_tag = self.workflow_task.task_tag
+        self.task_id = self.workflow_task.task_id
+        self.episode_uuid = self.workflow_task.episode_uuid
+
         self.config = config
         self.tokenizer = tokenizer
         self.saved_timelines: List[List[ExtendedMessage]] = []
@@ -136,7 +141,7 @@ def __init__(self, config, tokenizer, **kwargs):
         self.already_mad_flag: bool = False
         self.round_cnt = 0
         self.generation_prompt_token = None
-        self.workflow_metadata: Optional[Dict[str, Any]] = None  # Initialize workflow_metadata to store tool statistics
+        self.log_metrics: Optional[Dict[str, Union[float, List[float]]]] = None  # Initialize workflow_metadata to store tool statistics
 
         assert (
             self.config.ajet.data.max_prompt_length

ajet/context_tracker/multiagent_tracking.py

@@ -50,7 +50,6 @@ def __init__(
         config,
         should_interrupt_fn,
         generated_token_callback_fn,
-        episode_uuid: str,
         **kwargs,
     ):
         super().__init__(config, tokenizer, **kwargs)
@@ -61,7 +60,6 @@ def __init__(
         self.output_kwargs = {}
         self.input_kwargs = {}
         self.timeline_cache = {}
-        self.episode_uuid = episode_uuid
 
 
     def preprocess_tools_field(self, tools: List[dict] = [], disable_toolcalls: bool = False):

ajet/default_config/ajet_default.yaml

@@ -231,30 +231,57 @@ ajet:
 
   # trainer common configurations
   trainer_common:
+
+    # validation before training
     val_before_train: False
     val_pass_n: 4
+
+    # save and test frequency (in step)
     save_freq: 20
     test_freq: 20
-    save_trajectory: False  # whether to save train/eval trajectories to JSON files
+
+    # total training epochs
     total_epochs: 50
+
     nnodes: 1
     n_gpus_per_node: 8
+
+    # logger selection
     logger: swanlab
+
+    # algorithm setting
     algorithm:
       adv_estimator: grpo
       use_kl_in_reward: False
+
+    # number of optimizer.step per big batch
     mini_batch_num: 1
+
+    # verl offload configs
     fsdp_config:
       param_offload: True
       optimizer_offload: True
+
+    # learning rate
     optim:
       lr: 1e-6
+
+    # enable KL loss regularization
     use_kl_loss: True
+
+    # kl divergence loss coefficient
     kl_loss_coef: 0.002
     kl_loss_type: low_var_kl
+
+    # Ulysses specific configs
     ulysses_sequence_parallel_size: 1
+
+    # base directory to save checkpoints
     checkpoint_base_dir: ./saved_checkpoints
 
+    # whether to save train/eval trajectories to JSON files
+    save_trajectory_as_json_file: False
+
 
 
 

ajet/launcher.py

@@ -16,7 +16,6 @@
 from ajet.utils.pty import pty_launch
 
 set_loguru_default_color()
-# load_dotenv()
 load_dotenv(override=False)
 
 

ajet/schema/task.py

@@ -43,3 +43,4 @@ class WorkflowOutput(BaseModel):
     reward: Union[float, List[float], None] = Field(default=None)
     is_success: Union[bool, None] = Field(default=None)
     metadata: Dict[str, Any] = Field(default_factory=dict)
+    log_metrics: Dict[str, Union[float, List[float]]] = Field(default_factory=dict)

ajet/task_rollout/resource_keeper.py

@@ -178,22 +178,23 @@ def step(self, action: dict) -> Tuple[str, float, bool, dict]:
         )
         obs = ""
         assert isinstance(env_output, dict)
-        # === Support list-type state passthrough ===
-        # 1. If state is a list (new standard format), pass through directly
+
         if isinstance(env_output["state"], list):
+            # 1. If state is a list (new standard format), pass through directly
             obs = env_output["state"]
-        # 2. If state is a dict (old format or error)
         else:
+            # 2. If state is a dict (old format or error)
             if ("content" not in env_output["state"]) and ("error" in env_output["state"]):
                 obs = f"[Error from environment: {env_output['error']}]"
             elif env_output["state"].get("content", "") == "":
                 obs = "Warning: the environment does not provide any feedback, please provide valid inpu and try again."
             else:
                 obs = env_output["state"]["content"]
+
         reward = 0
         info = {}
         terminate = env_output["is_terminated"]
-        return obs, reward, terminate, info
+        return obs, reward, terminate, info # type: ignore
 
     def reset(self) -> str:
         """Reset gym environment."""

ajet/task_runner/general_runner.py

@@ -17,9 +17,6 @@ class GeneralRunner(BaseAgentRunner):
     def execute(self, workflow_task: WorkflowTask) -> BaseContextTracker:
         observation_window = workflow_task.observation_window
         task_thread_index = workflow_task.task_thread_index
-        task_batch_index = workflow_task.task_batch_index
-        task_tag = workflow_task.task_tag
-        task_id = workflow_task.task_id
 
         workflow_import = self.config.ajet.rollout.user_workflow
         workflow_cls = dynamic_import(workflow_import)
@@ -34,10 +31,7 @@ def execute(self, workflow_task: WorkflowTask) -> BaseContextTracker:
             llm_inference_fn=self.llm_inference_fn,
             tokenizer=self.tokenizer,
             config=self.config,
-            task_batch_index=task_batch_index,
-            task_tag=task_tag,
-            task_id=task_id,
-            episode_uuid=workflow_task.episode_uuid,
+            workflow_task = workflow_task,
             **hooks,
         )
         tuner = AjetTuner(
@@ -46,38 +40,17 @@ def execute(self, workflow_task: WorkflowTask) -> BaseContextTracker:
             user_workflow=user_workflow,
             config=self.config,
         )
-
         workflow_output: WorkflowOutput = asyncio.run(
             user_workflow.execute(workflow_task, tuner)
         )
-        # set workflow metadata to context tracker metadata
-        context_tracker.workflow_metadata = workflow_output.metadata
         if workflow_output.reward is not None:
             raw_reward, is_success = (
                 workflow_output.reward,
                 workflow_output.is_success,
             )
         else:
             raw_reward, is_success = self.get_judge().compute_reward(workflow_task, workflow_output)
-            
-            # ✅ Critical Fix: After calling `judge`, write the updated `reward_stats` back to `workflow_metadata`
-            # # Ensure that `native_compat_trainer` reads the actual value calculated by `judge`, not the 0 value returned by `env`.
-            if workflow_output.metadata and 'reward_stats' in workflow_output.metadata:
-                context_tracker.workflow_metadata['reward_stats'] = workflow_output.metadata['reward_stats']
-            else:
-                # fallback: If the judge does not update reward_stats, use the default value.
-                logger.warning(f"[WARN] reward_stats not found in metadata after judge call, creating default values")
-                default_reward_stats = {
-                    'original_reward': raw_reward,  
-                    'penalty': 0.0,
-                    'step_reward': raw_reward,
-                }
-                if workflow_output.metadata:
-                    workflow_output.metadata['reward_stats'] = default_reward_stats
-                    context_tracker.workflow_metadata['reward_stats'] = default_reward_stats
-                else:
-                    context_tracker.workflow_metadata = {'reward_stats': default_reward_stats}
-        
+
         workflow_task.gym_env = None  # clear gym env client reference to avoid serialization issue
 
         assert not isinstance(
@@ -95,12 +68,11 @@ def execute(self, workflow_task: WorkflowTask) -> BaseContextTracker:
         )
         context_tracker.process_reward(reward)
         # generate token before merging
-        context_tracker.task_id = task_id
-        context_tracker.task_tag = task_tag
         context_tracker.group_merge()
         # after merging, process and align reward again
         context_tracker.process_reward(reward)
         # mark the thread as ended
         observation_window["step"][task_thread_index] = -1
         tuner.terminate_episode()
+        context_tracker.log_metrics = workflow_output.log_metrics
         return context_tracker

ajet/utils/metric_helper/__init__.py

@@ -0,0 +1,17 @@
+from ajet.utils.metric_helper.save_trajectory_as_json import save_trajectory_as_json
+from ajet.utils.metric_helper.tool_metric_helper import compute_tool_metrics_from_trajectories
+from ajet.utils.metric_helper.reward_metric_helper import compute_reward_metrics_from_trajectories
+
+
+def save_trajectory_as_json_file(ctx_trackers, global_steps, config, prefix):
+    if config.ajet.trainer_common.save_trajectory_as_json:
+        save_trajectory_as_json(ctx_trackers, global_steps, prefix)
+
+def update_metrics(context_tracker_arr, metrics:dict):
+    tool_metrics = compute_tool_metrics_from_trajectories(context_tracker_arr)
+    reward_metrics = compute_reward_metrics_from_trajectories(context_tracker_arr)
+    if tool_metrics:
+        metrics.update(tool_metrics)
+    if reward_metrics:
+        metrics.update(reward_metrics)
+    return

ajet/utils/metric_helper/save_trajectory_as_json.py

@@ -0,0 +1,56 @@
+import os
+import json
+from ajet.utils.msg_converter import convert_grouped_steps_to_openai_format
+
+
+def save_trajectory_as_json(ctx_trackers, global_steps, prefix="train"):
+    """
+    Save ctx_trackers to JSON files for either training or evaluation.
+
+    Args:
+        ctx_trackers (list): List of context trackers containing trajectory data.
+        global_steps (int): The global step count to organize saved files.
+        prefix (str): Directory prefix indicating the type of trajectory ("train" or "eval").
+    """
+    for ctx_tracker in ctx_trackers:
+        # Determine task tag based on reward
+        reward = ctx_tracker.reward_structure.raw_reward
+        if reward >= 1:
+            ctx_tracker.tag = "success"
+        elif reward == 0:
+            ctx_tracker.tag = "failure"
+        else:
+            ctx_tracker.tag = "half_success"
+
+        formatted_traj = convert_grouped_steps_to_openai_format(ctx_tracker.timeline_cache)
+
+        # Prepare trajectory data
+        traj_data = {
+            "task_id": ctx_tracker.task_id,
+            "task_tag": ctx_tracker.tag,
+            "reward_structure": ctx_tracker.reward_structure.model_dump(),
+            "traj": formatted_traj
+        }
+
+        # Extract reward_stats from workflow_metadata if available
+        if hasattr(ctx_tracker, 'workflow_metadata') and ctx_tracker.workflow_metadata:
+            if 'reward_stats' in ctx_tracker.workflow_metadata:
+                traj_data['reward_structure']['reward_stats'] = ctx_tracker.workflow_metadata['reward_stats']
+
+        # Define save directory and file path
+        traj_save_dir = os.path.join(
+            os.environ.get("BEST_LOGGER_PATH", "launcher_record"),
+            "ctx_trackers",
+            prefix,
+            f"step_{global_steps}"
+        )
+        os.makedirs(traj_save_dir, exist_ok=True)
+        traj_file_path = os.path.join(traj_save_dir, f"{ctx_tracker.task_id}.json")
+
+        # Save trajectory data to JSON file
+        with open(traj_file_path, "w", encoding="utf-8") as f:
+            json.dump(traj_data, f, ensure_ascii=False, indent=2)
+
+        # Print confirmation for evaluation trajectories
+        if prefix != "train":
+            print(f"Saved trajectory to {traj_file_path}")