feat: Enable simulated user for multi-turn GRPO

examples/configs/grpo_adk_llama8b.yaml

@@ -0,0 +1,44 @@
+# GRPO configuration for unique numbers environment
+defaults: "grpo_math_8B.yaml"
+
+grpo:
+  num_prompts_per_step: 32
+  num_generations_per_prompt: 16
+  max_rollout_turns: 20
+  max_num_steps: 100
+  val_at_start: false
+
+data:
+  add_system_prompt: false
+
+checkpointing:
+  enabled: false
+  checkpoint_dir: "results/grpo-adk"
+  metric_name: "val_reward"
+  higher_is_better: true
+  keep_top_k: 3
+  save_period: 10
+
+env:
+  unique_numbers:
+    cfg:
+      max_turns: 15
+      min_length: 5
+      max_length: 10
+      max_integer: 15
+
+logger:
+  wandb_enabled: True 
+  wandb:
+    project: "grpo-simulated-adk"
+    name: "llama-8b-__NOW__"
+
+policy:
+  train_global_batch_size: 512
+  dynamic_batching:
+    enabled: False
+  tokenizer:
+    chat_template: "{% for message in messages %}{% if loop.first %}<|begin_of_text|>{% endif %}<|start_header_id|>{{ message['role'] }}<|end_header_id|>\n{{ message['content'] }}<|eot_id|>{% endfor %}<|start_header_id|>assistant<|end_header_id|>\n"
+
+cluster:
+  gpus_per_node: 8

examples/run_grpo_unique_numbers_w_adk.py

@@ -0,0 +1,245 @@
+"""Run GRPO with the Unique Numbers Simulator using ADK.
+
+This script sets up and executes the Group Relative Policy Optimization (GRPO) algorithm
+in a multi-turn conversational environment powered by the ADK framework.
+
+### Task Overview
+The objective is to train an agent to guess the number of unique integers in a list generated by a simulated user.
+The interaction is structured as a turn-based dialogue:
+- The user generates a list of integers.
+- The agent queries specific positions in the list (by index).
+- The user replies with the value at that index (if available).
+- The agent continues the interaction until it makes a final guess at the number of unique integers.
+
+### Environment Details
+The environment is a simulated user that:
+- Randomly generates a list of integers at setup.
+- Responds to the agent's queries using an LLM via the ADK endpoint.
+- Optionally evaluates the agent's final guess using an LLM-based grader (included for extensibility, though not essential for this task).
+
+### Example Usage
+    uv run python examples/run_grpo_unique_numbers_w_adk.py
+
+### Requirements
+- A working ADK environment with access to a compatible LLM endpoint.
+  For the default Gemini endpoint, the following environment variables must be set:
+    - `GOOGLE_GENAI_USE_VERTEXAI=1`
+    - `GOOGLE_CLOUD_PROJECT="your-project-id"`
+    - `GOOGLE_CLOUD_LOCATION="your-location"`
+
+- A properly configured GRPO YAML file.
+  By default, the script uses:
+    `examples/configs/grpo_adk_llama8b.yaml`
+"""
+
+import argparse
+import itertools
+import os
+import pprint
+import random
+from datetime import datetime, timedelta
+from typing import Iterator
+
+from omegaconf import OmegaConf
+from torch.utils.data import IterableDataset
+from transformers import AutoTokenizer
+
+from nemo_rl.algorithms.grpo import MasterConfig, grpo_train, setup
+from nemo_rl.algorithms.utils import get_tokenizer
+from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType
+from nemo_rl.distributed.virtual_cluster import init_ray
+from nemo_rl.environments.simulated_user.prompt import starting_user_prompt
+from nemo_rl.environments.simulated_user.unique_numbers import (
+    UniqueNumbersEnv,
+    UniqueNumbersMetadata,
+)
+from nemo_rl.models.generation import configure_generation_config
+from nemo_rl.utils.config import load_config, parse_hydra_overrides
+from nemo_rl.utils.logger import get_next_experiment_dir
+
+OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run GRPO with unique numbers simulator"
+    )
+    parser.add_argument(
+        "--config", type=str, default=None, help="Path to YAML config file"
+    )
+    args, overrides = parser.parse_known_args()
+    return args, overrides
+
+
+def generate_datum(
+    tokenizer: AutoTokenizer,
+    env_cfg: dict,
+    task_name: str,
+    idx: int,
+    add_system_prompt: bool,
+) -> DatumSpec:
+    formatted_prompt = tokenizer.apply_chat_template(
+        [{"role": "user", "content": starting_user_prompt}],
+        tokenize=False,
+        add_system_prompt=add_system_prompt,
+        add_generation_prompt=True,
+        add_special_tokens=False,
+    ).strip()
+    token_ids = tokenizer(
+        formatted_prompt, return_tensors="pt", add_special_tokens=False
+    )["input_ids"][0]
+
+    def _generate_numbers(
+        min_length, max_length, max_integer, default_max_turns
+    ) -> UniqueNumbersMetadata:
+        length = random.randint(min_length, max_length)
+        numbers = [random.randint(0, max_integer) for _ in range(length)]
+        return UniqueNumbersMetadata(
+            numbers=numbers,
+            unique_count=len(set(numbers)),
+            turn=0,
+            max_turns=default_max_turns,
+        )
+
+    metadata = _generate_numbers(
+        min_length=env_cfg["cfg"]["min_length"],
+        max_length=env_cfg["cfg"]["max_length"],
+        max_integer=env_cfg["cfg"]["max_integer"],
+        default_max_turns=env_cfg["cfg"]["max_turns"],
+    )
+
+    message_log: LLMMessageLogType = [
+        {"role": "user", "content": formatted_prompt, "token_ids": token_ids}
+    ]
+    return {
+        "message_log": message_log,
+        "length": len(token_ids),
+        "extra_env_info": metadata,
+        "loss_multiplier": 1.0,
+        "idx": idx,
+        "task_name": task_name,
+    }
+
+
+class IterableNumbersDataset(IterableDataset):
+    def __init__(self, tokenizer, env_cfg, task_name, add_system_prompt, length):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.env_cfg = env_cfg
+        self.task_name = task_name
+        self.add_system_prompt = add_system_prompt
+        self.length = length
+
+    def __iter__(self) -> Iterator[DatumSpec]:
+        for i in itertools.count():
+            yield generate_datum(
+                tokenizer=self.tokenizer,
+                env_cfg=self.env_cfg,
+                task_name=self.task_name,
+                idx=i,
+                add_system_prompt=self.add_system_prompt,
+            )
+
+    def __len__(self):
+        return self.length
+
+
+def setup_data(tokenizer, env_cfg, task_name, length, val_length, add_system_prompt):
+    env_config = env_cfg[task_name]
+    env = UniqueNumbersEnv.options(num_gpus=0).remote(cfg=dict(env_config["cfg"]))
+    task_to_env = {task_name: env}
+
+    train_ds = IterableNumbersDataset(
+        tokenizer=tokenizer,
+        env_cfg=env_config,
+        task_name=task_name,
+        add_system_prompt=add_system_prompt,
+        length=length,
+    )
+    val_ds = IterableNumbersDataset(
+        tokenizer=tokenizer,
+        env_cfg=env_config,
+        task_name=task_name,
+        add_system_prompt=add_system_prompt,
+        length=val_length,
+    )
+    val_task_to_env = task_to_env
+    return train_ds, val_ds, task_to_env, val_task_to_env
+
+
+def main():
+    args, overrides = parse_args()
+    if not args.config:
+        args.config = os.path.join(
+            os.path.dirname(__file__), "configs", "grpo_adk_llama8b.yaml"
+        )
+    config = load_config(args.config)
+    if overrides:
+        config = parse_hydra_overrides(config, overrides)
+    config: MasterConfig = OmegaConf.to_container(config, resolve=True)
+
+    now_pst = datetime.utcnow() + timedelta(hours=-7)
+    config["logger"]["wandb"]["name"] = config["logger"]["wandb"]["name"].replace(
+        "__NOW__", now_pst.strftime("%m/%d-%H:%M")
+    )
+
+    config["logger"]["log_dir"] = get_next_experiment_dir(config["logger"]["log_dir"])
+    if config["checkpointing"]["enabled"]:
+        print(
+            f"\U0001f4ca Using checkpoint directory: {config['checkpointing']['checkpoint_dir']}"
+        )
+
+    pprint.pprint(config)
+
+    init_ray()
+
+    tokenizer = get_tokenizer(config["policy"]["tokenizer"])
+    config["policy"]["generation"] = configure_generation_config(
+        config["policy"]["generation"], tokenizer
+    )
+
+    ds_length = (
+        config["grpo"]["num_prompts_per_step"]
+        * config["grpo"]["num_generations_per_prompt"]
+        * config["grpo"]["max_num_steps"]
+    )
+    dataset, val_dataset, task_to_env, val_task_to_env = setup_data(
+        tokenizer=tokenizer,
+        env_cfg=config["env"],
+        task_name="unique_numbers",
+        length=ds_length,
+        val_length=config["grpo"]["max_val_samples"],
+        add_system_prompt=config["data"]["add_system_prompt"],
+    )
+
+    (
+        policy,
+        policy_generation,
+        cluster,
+        dataloader,
+        val_dataloader,
+        loss_fn,
+        logger,
+        checkpointer,
+        grpo_state,
+        master_config,
+    ) = setup(config, tokenizer, dataset, val_dataset)
+
+    grpo_train(
+        policy,
+        policy_generation,
+        dataloader,
+        val_dataloader,
+        tokenizer,
+        loss_fn,
+        task_to_env,
+        val_task_to_env,
+        logger,
+        checkpointer,
+        grpo_state,
+        master_config,
+    )
+
+
+if __name__ == "__main__":
+    main()

nemo_rl/distributed/ray_actor_environment_registry.py

@@ -22,6 +22,7 @@
     "nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker": PY_EXECUTABLES.MCORE,
     "nemo_rl.environments.math_environment.MathEnvironment": PY_EXECUTABLES.SYSTEM,
     "nemo_rl.environments.games.sliding_puzzle.SlidingPuzzleEnv": PY_EXECUTABLES.SYSTEM,
+    "nemo_rl.environments.simulated_user.unique_numbers.UniqueNumbersEnv": PY_EXECUTABLES.SYSTEM,
 }