Fix remote environments & episode-driven pipeline

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Author: tdene

megatron/rl/agent/reward_only_agent.py

@@ -239,7 +239,16 @@ async def prepare_group_rollout(
     ) -> GroupRolloutParams:
 
         prompt, golden = await self.get_prompt(validation=request.validation)
+        return self.group_rollout_params(request, prompt=prompt, golden=golden)
 
+    def group_rollout_params(
+        self,
+        request: RolloutRequest | GroupedRolloutRequest,
+        *,
+        prompt: str | list[LLMChatMessage],
+        golden: Any,
+    ) -> GroupRolloutParams:
+        """The params for one group's rollouts on an already sampled prompt."""
         # Every rollout runs as a (possibly multi-turn) episode over the group's shared prompt.
         return GroupRolloutParams(
             run_episode=functools.partial(self._run_episode, request, prompt=prompt, golden=golden),

megatron/rl/server/agent/fastapi_env_server.py

@@ -2,11 +2,13 @@
 
 import asyncio
 import socket
+from typing import Annotated, Any
 
 import httpx
 import yaml
 from fastapi import FastAPI
-from pydantic import BaseModel, Field, PrivateAttr
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel, Field, PrivateAttr, TypeAdapter
 from typing_extensions import Self
 from uvicorn import Config, Server
 from uvicorn.config import LOGGING_CONFIG
@@ -22,14 +24,25 @@
     EvaluationResponse,
     GroupedRolloutRequest,
     GroupRolloutParams,
-    RolloutGenerator,
+    Rollout,
     RolloutRequest,
     TokenRollout,
 )
 from ...agent.registry import get_agent_class
+from ...agent.reward_only_agent import RewardOnlyAgent
 from ...inference import InferenceInterface
 from ...inference.chat_interface import MegatronChatInterface
-from ..api import EnvironmentServer, RemoteEvaluationRequest, RemoteRolloutRequest
+from ..api import (
+    EnvironmentServer,
+    RemoteEpisodeRequest,
+    RemoteEvaluationRequest,
+    RemotePrompt,
+    RemotePromptRequest,
+    RemoteRolloutRequest,
+)
+
+# A finished rollout on the wire; an empty placeholder validates as the first member.
+_ROLLOUT = TypeAdapter(Annotated[TokenRollout | Rollout, Field(union_mode="left_to_right")])
 
 
 def _bind(wire: BaseModel, request_cls: type, interface: InferenceInterface):
@@ -45,30 +58,49 @@ class FastAPIEnvServer(EnvironmentServer):
 
     @classmethod
     def build_app(cls, env_cls: type[Agent], cls_args: dict, policy: InferenceInterface) -> FastAPI:
-        """Env server app; every route generates through `policy`."""
+        """Env server app: one environment instance generating through `policy`."""
         app = FastAPI()
-
-        if issubclass(env_cls, ContrastiveRolloutGenerator):
+        env = env_cls(**cls_args)
+
+        @app.exception_handler(Exception)
+        async def env_error(_, exc: Exception) -> JSONResponse:
+            # The trainer sees the environment's failure the way a local agent would raise it.
+            return JSONResponse(status_code=500, content={"detail": repr(exc)})
+
+        if isinstance(env, RewardOnlyAgent):
+
+            @app.post("/prompts/")
+            async def sample_prompt(request: RemotePromptRequest) -> RemotePrompt:
+                prompt, golden = await env.get_prompt(validation=request.validation)
+                return RemotePrompt(prompt=prompt, golden=golden)
+
+            @app.post("/episodes/")
+            async def run_episode(request: RemoteEpisodeRequest) -> dict[str, Any]:
+                group_request = GroupedRolloutRequest(
+                    num_groups=1,
+                    rollouts_per_group=1,
+                    inference_interface=policy,
+                    generation_args=request.generation_args,
+                    validation=request.validation,
+                )
+                params = env.group_rollout_params(
+                    group_request, prompt=request.prompt, golden=request.golden
+                )
+                rollout = await params.build_rollout(await params.run_episode())
+                return rollout.model_dump()
+
+        if isinstance(env, ContrastiveRolloutGenerator):
 
             @app.post("/contrastive_rollouts/")
             async def contrastive_rollouts(
                 request: RemoteRolloutRequest,
             ) -> list[ContrastiveRollout]:
-                env = env_cls(**cls_args)
                 return await env.get_contrastive_rollouts(_bind(request, RolloutRequest, policy))
 
-        if issubclass(env_cls, RolloutGenerator):
-
-            @app.post("/rollouts/")
-            async def rollouts(request: RemoteRolloutRequest) -> list[TokenRollout]:
-                env = env_cls(**cls_args)
-                return await env.get_reward_rollouts(_bind(request, RolloutRequest, policy))
-
-        if issubclass(env_cls, EvaluationAgent):
+        if isinstance(env, EvaluationAgent):
 
             @app.post("/evaluation/")
             async def run_evaluation(request: RemoteEvaluationRequest):
-                env = env_cls(**cls_args)
                 return await env.run_evaluation(_bind(request, EvaluationRequest, policy))
 
         return app
@@ -115,23 +147,39 @@ async def get_contrastive_rollouts(self, request: RolloutRequest) -> list[Contra
         result = await self._post("/contrastive_rollouts/", wire)
         return [ContrastiveRollout.model_validate(r) for r in result]
 
-    async def prepare_group_rollout(self, request: GroupedRolloutRequest) -> GroupRolloutParams:
-        raise NotImplementedError(
-            "FastAPIEnvServer overrides get_reward_rollouts; prepare_group_rollout is not used."
+    async def prepare_group_rollout(
+        self, request: GroupedRolloutRequest | RolloutRequest
+    ) -> GroupRolloutParams:
+        sampled = RemotePrompt.model_validate(
+            await self._post("/prompts/", RemotePromptRequest(validation=request.validation))
+        )
+        episode_request = RemoteEpisodeRequest(
+            prompt=sampled.prompt,
+            golden=sampled.golden,
+            generation_args=request.generation_args,
+            validation=request.validation,
         )
 
+        async def run_episode() -> Rollout | TokenRollout:
+            # The env server scores the episode itself, so the episode is already the rollout.
+            return _ROLLOUT.validate_python(await self._post("/episodes/", episode_request))
+
+        async def build_rollout(rollout: Rollout | TokenRollout) -> Rollout | TokenRollout:
+            return rollout
+
+        return GroupRolloutParams(run_episode=run_episode, build_rollout=build_rollout)
+
     async def get_rollout_response(self, request, inference_request):
         raise NotImplementedError(
-            "FastAPIEnvServer overrides get_reward_rollouts/run_evaluation; "
-            "get_rollout_response is not used."
+            "Episodes run on the env server; FastAPIEnvServer never generates itself."
         )
 
-    async def get_reward_rollouts(self, request: RolloutRequest) -> list[TokenRollout]:
-        wire = RemoteRolloutRequest.model_validate(
-            request.model_dump(exclude={"inference_interface"})
-        )
-        result = await self._post("/rollouts/", wire)
-        return [TokenRollout.model_validate(r) for r in result]
+    async def get_reward_rollouts(self, request: RolloutRequest) -> list[Rollout | TokenRollout]:
+        async def single_rollout() -> Rollout | TokenRollout:
+            params = await self.prepare_group_rollout(request)
+            return await params.build_rollout(await params.run_episode())
+
+        return list(await asyncio.gather(*[single_rollout() for _ in range(request.num_rollouts)]))
 
     async def run_evaluation(self, request: EvaluationRequest) -> EvaluationResponse:
         wire = RemoteEvaluationRequest.model_validate(

megatron/rl/server/api.py

@@ -1,11 +1,15 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+from typing import Any
+
 from pydantic import BaseModel, Field
 from typing_extensions import Self, Type
 
 from .. import TypeLookupable
+
+from ..__init__ import Request
 from ..agent.api import EvaluationRequest, RolloutRequest
-from ..inference import InferenceInterface
+from ..inference import InferenceInterface, LLMChatMessage
 
 
 class Server(TypeLookupable):
@@ -37,8 +41,31 @@ class EnvironmentServer(Server):
     ...
 
 
+class RemotePromptRequest(BaseModel):
+    """Trainer -> env server: sample the prompt one group of rollouts will share."""
+
+    validation: bool = False
+
+
+class RemotePrompt(BaseModel):
+    """Env server -> trainer: a sampled prompt and the golden data that scores it."""
+
+    prompt: str | list[LLMChatMessage]
+    golden: Any = None
+
+
 ### Intentionally force `inference_interface` to be `None` in these subclasses.
 
+
+class RemoteEpisodeRequest(Request):
+    """Trainer -> env server: run one episode on a sampled prompt and score it."""
+
+    prompt: str | list[LLMChatMessage]
+    golden: Any = None
+    validation: bool = False
+    inference_interface: None = None
+
+
 class RemoteRolloutRequest(RolloutRequest):
     inference_interface: None = None