Add refmodel stub

Author: joecummings

apps/grpo/main.py

@@ -19,8 +19,25 @@ async def update_weights(self):
         pass
 
 
-async def generate_rollout():
-    pass
+class Episode:
+
+    turns = []
+
+    def add_turn(self, turn):
+        self.turns.append(turn)
+
+    def add_transform_info(self, key, data):
+        setattr(self, key, data)
+
+
+class ComputeAdvantages(Actor):
+    def __call__(self, episode):
+        pass
+
+
+class RefModel(Actor):
+    def forward(self, x):
+        pass
 
 
 async def main():
@@ -46,11 +63,25 @@ async def main():
         batch_size=4,
         max_policy_age=1,
     )
+    dataloader = await spawn_service(
+        default_service_cfg,
+        ForgeDataset,
+        path="gsm8k",
+    )
 
     async def continuous_rollouts():
         while True:
+            prompt = await dataloader.__next__.call()
+            if prompt is None:
+                print(f"Dataloader is empty, exiting rollout creation")
+                return
             version = await policy.get_current_version.choose()
-            episode = await generate_rollout(version)
+            episode = Episode()
+            with policy.session(version=version):
+                action = await policy.generate.call(prompt)
+                episode.add_turn((prompt, action))
+            episode.add_advantages(await compute_advantages.__call__.call(episode))
+            episode.add_logprobs(await ref_model.forward.call(episode.get_tokens()))
             await replay_buffer.add.call(episode)
 
     rollout_task = asyncio.create_task(continuous_rollouts())