mv dapo workflow to boxed workflo

Author: hiyuchang

examples/dapo_math/dapo.yaml

@@ -32,6 +32,7 @@ buffer:
         logprobs: 0
       workflow_args:
         use_base: true
+      reward_fn_args:
         enable_overlong_penalty: true
         penalty_factor: 1.0
         max_response_length: 20480
@@ -49,7 +50,8 @@ buffer:
         n: 32
         temperature: 1.0
         top_p: 0.7
-    default_workflow_type: 'math_dapo_workflow'
+    default_workflow_type: 'math_boxed_workflow'
+    default_reward_fn_type: 'math_dapo_reward'
   trainer_input:
     experience_buffer:
       name: math_buffer

trinity/common/rewards/__init__.py

@@ -6,6 +6,7 @@
 
 from .accuracy_reward import AccuracyReward
 from .countdown_reward import CountDownRewardFn
+from .dapo_reward import MathDAPORewardFn
 from .format_reward import FormatReward
 from .math_reward import MathBoxedRewardFn, MathRewardFn
 

trinity/common/rewards/dapo_reward.py

@@ -32,6 +32,7 @@ def __call__(  # type: ignore
         response: str,
         response_token: torch.Tensor,
         truth: Optional[str] = None,
+        **kwargs,
     ) -> Union[float, dict]:
         accuracy_score = compute_score(response, truth)
 

trinity/common/rewards/math_reward.py

@@ -49,16 +49,17 @@ class MathBoxedRewardFn(RewardFn):
 
     def __init__(
         self,
+        **kwargs,
     ) -> None:
         pass
 
     def __call__(  # type: ignore
         self,
         response: str,
-        prompt: Optional[str] = None,
         truth: Optional[str] = None,
         with_think: Optional[bool] = False,
         format_score_coef: Optional[float] = 0.1,
+        **kwargs,
     ) -> dict[str, float]:
         accuracy_score = compute_score(response, truth)
 

trinity/common/workflows/__init__.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 """Workflow module"""
 from .customized_math_workflows import MathBoxedWorkflow
-from .dapo_workflow import MathDAPOWorkflow
 from .envs.alfworld.alfworld_workflow import AlfworldWorkflow
 from .envs.sciworld.sciworld_workflow import SciWorldWorkflow
 from .envs.webshop.webshop_workflow import WebShopWorkflow
@@ -19,5 +18,4 @@
     "SciWorldWorkflow",
     "MathBoxedWorkflow",
     "MathRMWorkflow",
-    "MathDAPOWorkflow",
 ]

trinity/common/workflows/customized_math_workflows.py

@@ -31,6 +31,7 @@ def reset(self, task: Task):
         self.is_eval = task.is_eval
 
         self.workflow_args = task.workflow_args
+        self.reward_fn_args = task.reward_fn_args
 
         self.use_base = self.workflow_args.get("use_base", False)
         self.with_think = self.workflow_args.get("with_think", False)
@@ -49,9 +50,18 @@ def reset(self, task: Task):
                 self.system_prompt = default_prompt
 
         if task.reward_fn is None:
-            self.reward_fn = MathBoxedRewardFn()
+            self.reward_fn = MathBoxedRewardFn(**self.reward_fn_args)
         else:
-            self.reward_fn = task.reward_fn
+            self.reward_fn = task.reward_fn(**self.reward_fn_args)
+
+    def format_prompt(self):
+        prompt_text = ""
+        if self.system_prompt:
+            prompt_text += "System:" + self.system_prompt
+            prompt_text += "\nUser:\n" + self.task_desc + "\nAssistant:\n"
+        else:
+            prompt_text += "User:\n" + self.task_desc + "\nAssistant:\n"
+        return prompt_text
 
     def run(self) -> List[Experience]:
         if not self.use_base:
@@ -71,6 +81,7 @@ def run(self) -> List[Experience]:
                 truth=self.truth,
                 with_think=self.with_think,
                 format_score_coef=self.format_score_coef,
+                response_token=response.tokens[response.prompt_length :],
             )
 
             if response.metrics is None:
@@ -79,7 +90,12 @@ def run(self) -> List[Experience]:
             reward = sum(reward_dict.values())
             response.reward = reward
 
-            logger.debug(
-                f"self.task_desc: {self.task_desc}, messages: {messages}, response: {response.response_text}, reward: {reward}"
-            )
+            if not self.use_base:
+                logger.debug(
+                    f"self.task_desc: {self.task_desc}, messages: {messages}, response: {response.response_text}, reward: {reward}"
+                )
+            else:
+                logger.debug(
+                    f"self.task_desc: {self.task_desc}, prompt_text: {prompt_text}, response: {response.response_text}, reward: {reward}"
+                )
         return responses

trinity/common/workflows/dapo_workflow.py

@@ -206,16 +206,6 @@ def format_messages(self):
             messages.append({"role": "assistant", "content": self.reply_prefix})
         return messages
 
-    def format_prompt(self):
-        """Format prompt for the base model."""
-        prompt_text = ""
-        if self.system_prompt:
-            prompt_text += "System:\n" + self.system_prompt
-            prompt_text += "\nUser:\n" + self.task_desc + "\nAssistant:\n"
-        else:
-            prompt_text += "User:\n" + self.task_desc + "\nAssistant:\n"
-        return prompt_text
-
     def run(self) -> List[Experience]:
         # TODO: Optimize the generate function
         messages = self.format_messages()