Added is_correct & reward flow through tool env (#277)

* Added is_correct & reward flow through tool env

* Display rewards for all trajectories in episode, not just the first

Author: mananroongta

rllm/agents/tool_agent.py

@@ -94,6 +94,11 @@ def update_from_env(self, observation: Any, reward: float, done: bool, info: dic
         self.messages.extend(obs_messages)
         self.current_observation = observation
 
+        if self._trajectory.steps:
+            self._trajectory.steps[-1].reward = reward
+            self._trajectory.steps[-1].done = done
+            self._trajectory.steps[-1].info = info
+
     def update_from_model(self, response: str, **kwargs) -> Action:
         """
         Updates the agent's state based on the model's response.

rllm/engine/agent_workflow_engine.py

@@ -85,7 +85,9 @@ async def process_task_with_retry(self, task: dict, task_id: str, rollout_idx: i
                 uid = f"{task_id}:{rollout_idx}"
                 episode = await workflow.run_with_termination_handling(task=task, uid=uid, **kwargs)
 
-                colorful_print(f"[{uid}] Rollout completed with termination reason: {episode.termination_reason}", fg="green" if episode.is_correct else "yellow")
+                # Display rewards for all trajectories
+                rewards_str = ", ".join([f"{traj.name}: {traj.reward:.1f}" for traj in episode.trajectories])
+                colorful_print(f"[{uid}] Rollout completed. Rewards: {rewards_str}, Termination: {episode.termination_reason}", fg="green" if episode.is_correct else "yellow")
 
                 if episode.termination_reason != TerminationReason.ERROR:
                     return task_id, rollout_idx, episode

rllm/environments/tools/tool_env.py

@@ -98,7 +98,7 @@ def step(self, action: list[dict] | str | dict):
 
             task_info = self.task if self.task is not None else {}
             reward_output = self.reward_fn(task_info=task_info, action=llm_response)
-            return {}, reward_output.reward, done, {"response": action, "metadata": reward_output.metadata}
+            return {}, reward_output.reward, done, {"response": action, "metadata": reward_output.metadata, "is_correct": reward_output.is_correct}
 
         tool_calls = action
         assert isinstance(tool_calls, list)