feat: add val_only option and markdown output for validation metrics

- Add val_only config to run validation without training
- Add val_print_to_markdown_file_path to output validation metrics to file
- Rename TGC metrics to task_pass_rate and add task_pass_rate@k for k in [2,4,8,16]
- Add std_reward metric
- Add assertion to prevent incompatible swarm_mode with val_before_train/val_only

Author: binary-husky

ajet/backbone/__init__.py

@@ -13,4 +13,5 @@
         "AjetTaskReader",
     ]
 except ImportError:
-    logger.info("trinity is not available.")
+    pass
+    # logger.info("trinity is not available.")

ajet/backbone/trainer_verl.py

@@ -419,14 +419,23 @@ def fit(self):  # noqa: C901
         self.checkpoint_manager.update_weights(self.global_steps)
         self.checkpoint_manager.sleep_replicas()
 
+        # [oc] swarm_mode is not compatible with `val_before_train` and `val_only`
+        assert not (self.config.ajet.enable_swarm_mode and (self.config.ajet.trainer_common.val_before_train or self.config.ajet.trainer_common.val_only)), \
+            "swarm_mode is not compatible with `val_before_train` and `val_only`"
+
+
         # perform validation before training
         # currently, we only support validation using the reward_function.
-        if (self.val_reward_fn is not None) and (self.config.trainer.get("val_before_train", True)) and (not self.config.ajet.enable_swarm_mode):
+        if (self.val_reward_fn is not None) and (self.config.ajet.trainer_common.val_before_train) and (not self.config.ajet.enable_swarm_mode):
             val_metrics = self._validate()
             assert val_metrics, f"{val_metrics=}"
-            pprint(f"Initial validation metrics: {val_metrics}")
             self.verl_logger.log(data=val_metrics, step=self.global_steps)
-            if self.config.trainer.get("val_only", False):
+            val_print_to_markdown_file_path = self.config.ajet.trainer_common.val_print_to_markdown_file_path
+            if val_print_to_markdown_file_path:
+                with open(val_print_to_markdown_file_path, mode="a+") as f:
+                    f.write(str(val_metrics))
+                    f.write('\n')
+            if self.config.ajet.trainer_common.val_only:
                 return
 
         # add tqdm
@@ -983,11 +992,21 @@ def _rollout_val_dataset(self, target_dataset, target_dataset_name, mode, epoch)
             "total_tasks": len(task_results),
             "num_all_success_tasks": num_all_success_tasks,
             f"num_pass_n_tasks(pass@{pass_n})": num_pass_n_tasks,
-            "TGC@1": repeated_success_tasks / (num_tasks * pass_n),
-            f"TGC@{pass_n}": num_pass_n_tasks / num_tasks,
-            f"TGC@{pass_n}-all-pass": num_all_success_tasks / num_tasks,
+            # [oc]: change var name TGC -> task_pass_rate
+            "task_pass_rate@1": repeated_success_tasks / (num_tasks * pass_n),
+            f"task_pass_rate@{pass_n}": num_pass_n_tasks / num_tasks,
+            f"task_pass_rate@{pass_n}-all-pass": num_all_success_tasks / num_tasks,
             "mean_reward": sum(rewards) / len(rewards) if rewards else 0,
+            "std_reward": np.std(rewards) if rewards else 0,
         }
+        for k in [2, 4, 8, 16]:
+            if pass_n > k:
+                num_pass_k = 0
+                for task_id, task_outcomes in task_results.items():
+                    if any(tag == "success" for tag in task_outcomes["tag_arr"][:k]):
+                        num_pass_k += 1
+                val_metrics[f"task_pass_rate@{k}"] = num_pass_k / num_tasks
+
         save_trajectory_as_json_file(ctx_trackers, self.global_steps, self.config, prefix="eval")
         update_metrics(ctx_trackers, val_metrics, prefix="eval_")
         print_dict(

ajet/default_config/ajet_default.yaml

@@ -234,6 +234,8 @@ ajet:
     # validation before training
     val_before_train: False
     val_pass_n: 4
+    val_only: False
+    val_print_to_markdown_file_path: null
 
     # save and test frequency (in step)
     save_freq: 20

ajet/default_config/verl/config_auto_convertion_verl.jsonc

@@ -2,6 +2,7 @@
     "ajet.trainer_common.total_epochs": "trainer.total_epochs",
 
     "ajet.trainer_common.val_before_train": "trainer.val_before_train",
+    "ajet.trainer_common.val_only": "trainer.val_only",
     "ajet.trainer_common.n_gpus_per_node": "trainer.n_gpus_per_node",
     "ajet.trainer_common.nnodes": "trainer.nnodes",
     "ajet.trainer_common.logger": "trainer.logger",

tests/bench/benchmark_math/benchmark_math.yaml

@@ -43,14 +43,16 @@ ajet:
     debug_first_n_tasks: 1
 
   trainer_common:
-    save_freq: 100
+    save_freq: 999999
     test_freq: 999999
     total_epochs: 100
     logger: swanlab
+    val_print_to_markdown_file_path: './qwen2-7b.md'
     nnodes: 1
     n_gpus_per_node: 4
     # loss = loss * loss_extra_scale_ratio
     loss_extra_scale_ratio: 1.0
+    val_before_train: true
 
 
   execute_test: True # DO NOT EDIT, THIS IS FOR TEST ROBOT

tutorial/opencode_build_openclaw_interactive_train/fake_vllm_endpoint.py

@@ -299,7 +299,6 @@ async def finalize_episodes(task: Task, valid_results: List[EpisodeResult], rewa
             lambda ep=episode_result, wo=workflow_output: swarm_client.end_episode(task, ep.episode_uuid, wo),
         )
 
-# [oc]: 微调一下代码，当 handle_one2many_request_run_episodes 运行完时，随机stream回去一个答案，但stream不中断，等待reward计算完之后，再结束stream
 async def handle_one2many_request(request: Request, request_id: str) -> Dict | List[bytes]:
     task, valid_results, all_answers, user_query, all_answers = await handle_one2many_request_run_episodes(request, request_id)
     best_answer = await handle_one2many_request_run_reward(task, valid_results, all_answers, user_query)