Integrate metric logger into grpo training loop

apps/grpo/main.py

@@ -12,8 +12,10 @@
 import torch
 from datasets import load_dataset
 from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
+from forge.actors.replay_buffer import ReplayBuffer
 from forge.controller import ServiceConfig, spawn_service
 from forge.controller.actor import ForgeActor
+from forge.util.metric_logging import get_metric_logger
 from monarch.actor import endpoint
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -388,6 +390,13 @@ async def main():
     group_size = 1
     model = "Qwen/Qwen3-1.7B"
 
+    # ---- Setup WandB Logger ---- #
+    logger = get_metric_logger(
+        "wandb",
+        freq=1,
+        project="grpo-training",
+    )
+
     # ---- Setup services ---- #
     default_service_cfg = ServiceConfig(
         procs_per_replica=1,
@@ -498,6 +507,7 @@ async def continuous_rollouts():
                 print(
                     f"Generated {rollout_count} rollouts w/ average reward {avg_reward}"
                 )
+                logger.log("reward/rollout", avg_reward, rollout_count)
 
     async def continuous_training():
         training_step = 0
@@ -511,7 +521,9 @@ async def continuous_training():
                 if training_step % 10 == 0:
                     print(f"Completed {training_step} training steps")
                     if training_result:
-                        print(f"Latest loss: {training_result.get('loss', 'N/A')}")
+                        loss_value = training_result.get("loss", 0.0)
+                        print(f"Latest loss: {loss_value}")
+                        logger.log("loss/training_step", loss_value, training_step)
                 # await trainer.update_weights(policy)
 
     print("Starting GRPO training loops...")