Integrate Eval Protocol as RL environment (#276)

* fix issue #259

* init

* print(f"Train dataset size: {len(train_dataset)}")

* super().__init__(rollout_engine, **kwargs)

* qwen2p5-vl-32b-instruct

* comment

* print

* [trajectory]

* 100 tasks

* 10

* n_parallel_tasks = 10

* update to qwen3 and new server

* thinking

* add train_frozenlake_flow

* model_id = "accounts/pyroworks/deployedModels/qwen3-8b-g0m657sn"

* print(episode.trajectories)

* from datasets import load_dataset

* fixed concurrency

* tracing_signal

* signal

* trainer.val_before_train=False \

* _build_rollout_processor_config

* print("model in frozen_lake_flow", model)

* semaphore=self._rollout_processor_semaphore,

* print(error_message)

* update self.model  to deployedModel

* remove some codes

* formatting

---------

Co-authored-by: listar2000 <lisd.star2015@outlook.com>
Co-authored-by: Derek Xu <xzrderek@gmail.com>

Author: 3 people

examples/eval_protocol/frozen_lake_flow.py

@@ -0,0 +1,227 @@
+"""
+This workflow bridges eval-protocol's MCPGymRolloutProcessor with rllm-fw's Workflow pattern
+for the FrozenLake environment.
+"""
+
+import asyncio
+from pathlib import Path
+
+import eval_protocol
+from eval_protocol.benchmarks.test_frozen_lake import test_frozen_lake_evaluation
+from eval_protocol.models import EvaluationRow, InputMetadata, Message
+from eval_protocol.pytest.default_mcp_gym_rollout_processor import (
+    MCPGymRolloutProcessor,
+)
+from eval_protocol.pytest.types import RolloutProcessorConfig
+
+from rllm.agents.agent import Episode, Step, Trajectory
+from rllm.engine.rollout.openai_engine import OpenAIEngine
+from rllm.workflows.workflow import Workflow
+
+
+class FrozenLakeWorkflow(Workflow):
+    """
+    Workflow that executes frozen lake tasks using MCPGymRolloutProcessor.
+
+    Task format expected:
+    {
+        "id": "frozen_lake_task_0",
+        "system_prompt": "...",
+        "environment_context": {...},
+        "user_prompt_template": "{observation}"
+    }
+    """
+
+    # Class variables (shared across all workflow instances)
+    _shared_server_started = False
+    _server_lock = asyncio.Lock()
+    _shared_rollout_processor = MCPGymRolloutProcessor()
+
+    def __init__(self, rollout_engine: OpenAIEngine, lite_llm_prefix: str = "fireworks_ai/", max_steps: int = 30, temperature: float = 1.0, max_tokens: int = 4096, **kwargs):
+        super().__init__(rollout_engine, **kwargs)
+
+        self._rollout_processor_server_started = False
+        self._rollout_processor_semaphore = asyncio.Semaphore(1)
+        self._lite_llm_prefix = lite_llm_prefix
+        self._temperature = temperature
+        self._max_tokens = max_tokens
+        self._max_steps = max_steps
+
+        eval_protocol_path = Path(eval_protocol.__file__).parent
+        self._server_script_path = eval_protocol_path / "mcp_servers" / "frozen_lake" / "server.py"
+
+        # Use shared rollout processor across all instances
+        self.rollout_processor = FrozenLakeWorkflow._shared_rollout_processor
+
+    def _build_rollout_processor_config(self):
+        model = self._lite_llm_prefix + self.rollout_engine.model
+        print("model in frozen_lake_flow", model)
+        return RolloutProcessorConfig(
+            completion_params={
+                "model": model,
+                "temperature": self._temperature,
+                "max_tokens": self._max_tokens,
+            },
+            mcp_config_path="",
+            server_script_path=str(self._server_script_path),
+            steps=self._max_steps,
+            semaphore=self._rollout_processor_semaphore,
+            kwargs={"start_server": self._rollout_processor_server_started},
+        )
+
+    async def run(self, task: dict, uid: str, **kwargs) -> Episode:
+        """
+        Execute the frozen lake workflow.
+
+        Args:
+            task: Dict containing frozen lake task data
+            uid: Unique identifier for this episode
+            **kwargs: Additional arguments
+
+        Returns:
+            Episode with trajectory and computed rewards
+        """
+        # Thread-safe server startup (double-checked locking pattern)
+        if not FrozenLakeWorkflow._shared_server_started:
+            # Only acquire lock if server not started yet
+            async with FrozenLakeWorkflow._server_lock:
+                # Check again inside lock (another workflow might have started it)
+                if not FrozenLakeWorkflow._shared_server_started:
+                    # First workflow to reach here starts the server
+                    self._rollout_processor_server_started = True
+                    FrozenLakeWorkflow._shared_server_started = True
+                else:
+                    self._rollout_processor_server_started = False
+        else:
+            self._rollout_processor_server_started = False
+
+        self.reset(task=task, uid=uid)
+
+        try:
+            eval_row = self._task_to_evaluation_row(task)
+
+            tasks = self.rollout_processor([eval_row], self._build_rollout_processor_config())
+
+            if not tasks:
+                raise ValueError("MCPGymRolloutProcessor returned no tasks")
+
+            result_row: EvaluationRow = await tasks[0]
+
+            episode = await self._evaluate_and_create_episode(result_row, task, uid)
+
+            return episode
+
+        except Exception as e:
+            # Gracefully handle failures - return a failed Episode instead of crashing
+            print(f"⚠️  Task {uid} failed: {e}")
+
+            failed_episode = Episode(
+                id=uid,
+                task=task,
+                is_correct=False,
+                trajectories=[],
+                metrics={"frozen_lake_reward": 0.0, "error": str(e)},
+            )
+            return failed_episode
+
+    def _task_to_evaluation_row(self, task: dict) -> EvaluationRow:
+        """Convert rllm task dict to eval protocol EvaluationRow."""
+        return EvaluationRow(
+            messages=[Message(role="system", content=task["system_prompt"])],
+            input_metadata=InputMetadata(
+                row_id=task["id"],
+                dataset_info={
+                    "environment_context": task["environment_context"],
+                    "user_prompt_template": task["user_prompt_template"],
+                },
+            ),
+        )
+
+    async def _evaluate_and_create_episode(
+        self,
+        row: EvaluationRow,
+        task: dict,
+        uid: str,
+    ) -> Episode:
+        """
+        Evaluate the rollout and convert to rllm Episode.
+        """
+        # Call the evaluation function
+        evaluated_row: EvaluationRow = await test_frozen_lake_evaluation(row)
+
+        # Extract reward and metrics from evaluation_result
+        if evaluated_row.evaluation_result is None:
+            raise ValueError("Evaluation function did not return a result")
+
+        reward = evaluated_row.evaluation_result.score
+        reward_info = evaluated_row.evaluation_result.metrics or {}
+
+        def msg_to_dict(msg: Message) -> dict:
+            """Convert eval_protocol Message to chat completion dict."""
+            d = {"role": msg.role, "content": msg.content}
+            if msg.tool_calls:
+                d["tool_calls"] = [
+                    {
+                        "id": tc.id,
+                        "type": tc.type,
+                        "function": {
+                            "name": tc.function.name,
+                            "arguments": tc.function.arguments,
+                        },
+                    }
+                    for tc in msg.tool_calls
+                ]
+            if msg.tool_call_id:
+                d["tool_call_id"] = msg.tool_call_id
+            if msg.name:
+                d["name"] = msg.name
+            return d
+
+        trajectory = Trajectory()
+        all_messages = []
+
+        for msg in row.messages:
+            msg_dict = msg_to_dict(msg)
+            all_messages.append(msg_dict)
+
+            # Create Step with only observation and chat_completions for user or tool message
+            if msg.role in ["user", "tool"]:
+                new_step = Step(observation=str(msg.content or ""), chat_completions=all_messages.copy())
+                trajectory.steps.append(new_step)
+
+            # Create new Step with action/response for assistant message
+            elif msg.role == "assistant":
+                # Extract action: tool calls if present, otherwise message content
+                action_data = msg_dict.get("tool_calls") if msg.tool_calls else str(msg.content or "")
+
+                new_step = Step(
+                    model_response=str(msg.content) if msg.content else "",
+                    action=action_data,
+                    chat_completions=all_messages.copy(),
+                )
+                trajectory.steps.append(new_step)
+
+        # Assign final reward to the last step (sparse reward)
+        if trajectory.steps:
+            trajectory.steps[-1].reward = reward
+            trajectory.steps[-1].info = reward_info
+
+        trajectory.reward = reward
+        trajectory.task = task
+
+        # Create episode
+        episode = Episode(
+            id=uid,
+            task=task,
+            is_correct=(reward == 1.0),
+            trajectories=[trajectory],
+            metrics={"frozen_lake_reward": reward, **reward_info},
+        )
+
+        return episode
+
+    def cleanup(self):
+        """Cleanup MCP server resources."""
+        if self.rollout_processor:
+            self.rollout_processor.cleanup()
+            self.rollout_processor = None

examples/eval_protocol/prepare_frozen_lake_data.py

@@ -0,0 +1,33 @@
+import random
+
+from datasets import Dataset
+
+from rllm.data.dataset import DatasetRegistry
+
+
+def prepare_frozen_lake_data(train_size: int, test_size: int):
+    system_prompt = "You are playing FrozenLake, a grid-based navigation game displayed as a 4x4 text grid. The grid contains: S (Start), F (Frozen safe), H (Hole - deadly), G (Goal). You start at position S and must reach G while avoiding H tiles. In this version, the surface is not slippery so your moves are deterministic. IMPORTANT: When you are at the starting position, you appear as 'S'. When you move to other positions, the hightlighted position will change on the grid. If you step on H, the episode ends with failure.  Use the lake_move tool with actions LEFT, DOWN, RIGHT, UP to navigate the grid."
+    user_prompt_template = "Current game state grid:\n{observation}\n\nYou are navigating the 4x4 grid above. Navigate safely to reach the goal 'G' while avoiding holes 'H'. Choose your next move from: LEFT, DOWN, RIGHT, or UP."
+
+    def create_row(idx, seed):
+        return {"id": f"run_{idx}", "system_prompt": system_prompt, "user_prompt_template": user_prompt_template, "environment_context": {"game": "FrozenLake", "map_name": "4x4", "seed": seed}}
+
+    seeds = random.sample(range(1, 1_000_001), train_size + test_size)
+    all_rows = []
+    for i in range(train_size + test_size):
+        all_rows.append(create_row(i, seeds[i]))
+    train_rows = all_rows[:train_size]
+    test_rows = all_rows[train_size:]
+
+    train_dataset = Dataset.from_list(train_rows)
+    test_dataset = Dataset.from_list(test_rows)
+
+    DatasetRegistry.register_dataset("frozen_lake_eval_protocol", train_dataset, "train")
+    DatasetRegistry.register_dataset("frozen_lake_eval_protocol", test_dataset, "test")
+
+    print(f"Train dataset size: {len(train_dataset)}")
+    print(f"Test dataset size: {len(test_dataset)}")
+
+
+if __name__ == "__main__":
+    prepare_frozen_lake_data(train_size=100, test_size=100)

examples/eval_protocol/run_frozen_lake_flow.py

@@ -0,0 +1,118 @@
+"""
+Run Frozen Lake Workflow with rllm-fw
+
+This script demonstrates how to execute frozen lake tasks using rllm-fw's
+AgentWorkflowEngine with eval-protocol's MCPGymRolloutProcessor.
+"""
+
+import asyncio
+import json
+import os
+from pathlib import Path
+
+from frozen_lake_flow import FrozenLakeWorkflow
+
+from rllm.data.dataset import DatasetRegistry
+from rllm.engine.agent_workflow_engine import AgentWorkflowEngine
+from rllm.engine.rollout.openai_engine import OpenAIEngine
+
+
+def evaluate_results(episodes):
+    """
+    Evaluate the results and compute accuracy metrics.
+
+    Args:
+        episodes: List of Episode objects
+    """
+    total = len(episodes)
+    correct = sum(1 for ep in episodes if ep.is_correct)
+    accuracy = correct / total if total > 0 else 0.0
+
+    print("\n" + "=" * 60)
+    print("EVALUATION RESULTS")
+    print("=" * 60)
+    print(f"Total tasks: {total}")
+    print(f"Correct: {correct}")
+    print(f"Accuracy: {accuracy:.2%}")
+    print()
+
+    for episode in episodes:
+        status = "✅" if episode.is_correct else "❌"
+        reward = episode.metrics.get("frozen_lake_reward", 0.0)
+        print(f"{status} Task {episode.id}: reward={reward:.3f}")
+
+    print("=" * 60)
+
+    return accuracy
+
+
+async def main():
+    """Main execution function."""
+
+    n_parallel_tasks = 4
+    max_tasks = 4
+    model_id = "accounts/pyroworks/deployedModels/qwen3-8b-g0m657sn"
+
+    # Create dummy rollout_engine (required by Workflow base class but not used)
+    rollout_engine = OpenAIEngine(
+        model=model_id,
+        base_url="https://api.fireworks.ai/inference/v1",
+        api_key=os.getenv("FIREWORKS_API_KEY"),
+    )
+
+    engine = AgentWorkflowEngine(
+        workflow_cls=FrozenLakeWorkflow,
+        workflow_args={
+            "lite_llm_prefix": "fireworks_ai/",
+            "steps": 30,
+            "temperature": 1.0,
+            "max_tokens": 16384,
+        },
+        rollout_engine=rollout_engine,
+        n_parallel_tasks=n_parallel_tasks,
+        retry_limit=1,
+    )
+
+    test_dataset = DatasetRegistry.load_dataset("frozen_lake_eval_protocol", "test")
+    tasks = []
+    for i in range(max_tasks):
+        tasks.append(test_dataset[i])
+
+    print("Starting frozen lake workflow execution...")
+    print(f"Model: {model_id}")
+    print(f"Parallel tasks: {n_parallel_tasks}")
+    print()
+
+    try:
+        episodes = await engine.execute_tasks(tasks)
+        for episode in episodes:
+            print(episode.trajectories)
+        accuracy = evaluate_results(episodes)
+
+        output_dir = Path("logs")
+        output_dir.mkdir(exist_ok=True)
+        output_file = output_dir / "frozen_lake_results.json"
+
+        with open(output_file, "w") as f:
+            json.dump([episode.to_dict() for episode in episodes], f, indent=2)
+
+        print(f"\n✅ Results saved to {output_file}")
+
+        return accuracy
+
+    except Exception as e:
+        print(f"❌ Error during execution: {e}")
+        import traceback
+
+        traceback.print_exc()
+        raise
+    finally:
+        engine.shutdown()
+
+
+if __name__ == "__main__":
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    accuracy = asyncio.run(main())
+
+    print(f"\n🎯 Final Accuracy: {accuracy:.2%}")