feat: add trajectory saving feature with config control

- Add save_trajectory.py module with save_train_trajectory and save_eval_trajectory functions
- Add save_trajectory config flag in ajet_default.yaml (default: False)
- Integrate trajectory saving in trainer_verl.py for both training and evaluation
- Extract and save reward_structure, workflow_metadata, and OpenAI-formatted trajectories

Author: TaoShuchang

ajet/backbone/trainer_verl.py

@@ -54,6 +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
 
 
 def parse_reward_from_dataproto(data: DataProto, return_dict=False) -> dict | torch.Tensor:
@@ -577,6 +578,9 @@ 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")
@@ -1029,6 +1033,10 @@ 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

ajet/default_config/ajet_default.yaml

@@ -235,6 +235,7 @@ ajet:
     val_pass_n: 4
     save_freq: 20
     test_freq: 20
+    save_trajectory: False  # whether to save train/eval trajectories to JSON files
     total_epochs: 50
     nnodes: 1
     n_gpus_per_node: 8

ajet/utils/save_trajectory.py

@@ -0,0 +1,80 @@
+import os
+import json
+from ajet.utils.msg_converter import convert_grouped_steps_to_openai_format
+
+
+def save_train_trajectory(ctx_trackers, global_steps):
+    """Save training ctx_trackers to JSON files."""
+    for ctx_tracker in ctx_trackers:
+        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"
+        
+        # Use unified conversion function to convert grouped_steps to OpenAI format
+        if hasattr(ctx_tracker, 'get_grouped_steps_openai_format'):
+            formatted_traj = ctx_tracker.get_grouped_steps_openai_format()
+        else:
+            formatted_traj = convert_grouped_steps_to_openai_format(ctx_tracker.grouped_steps)
+
+        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 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']
+        
+        traj_save_dir = os.path.join(
+            os.environ.get("BEST_LOGGER_PATH", "launcher_record"),
+            "ctx_trackers",
+            "train",
+            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")
+        
+        with open(traj_file_path, "w", encoding="utf-8") as f:
+            json.dump(traj_data, f, ensure_ascii=False, indent=2)
+
+
+def save_eval_trajectory(ctx_trackers, global_steps):
+    """Save evaluation ctx_trackers to JSON files."""
+    for ctx_tracker in ctx_trackers:
+        # Use unified conversion function to convert grouped_steps to OpenAI format
+        if hasattr(ctx_tracker, 'get_grouped_steps_openai_format'):
+            formatted_traj = ctx_tracker.get_grouped_steps_openai_format()
+        else:
+            formatted_traj = convert_grouped_steps_to_openai_format(ctx_tracker.grouped_steps)
+
+        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 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']
+        
+        traj_save_dir = os.path.join(
+            os.environ.get("BEST_LOGGER_PATH", "launcher_record"),
+            "ctx_trackers",
+            "val",
+            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")
+        
+        with open(traj_file_path, "w", encoding="utf-8") as f:
+            json.dump(traj_data, f, ensure_ascii=False, indent=2)
+        
+        print(f"Saved trajectory to {traj_file_path}")