HRL cherry-pick: dev181 + fixes, skip broken #2099

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 Documenting changes which affect configuration usage patterns (added/moved/removed/renamed fields, notable logic changes).
 
+- **`orchestrator.env[].num_workers`**: Added configurable env server worker count (`int | "auto"`, default: `"auto"`). When `"auto"`, scales based on concurrency (1 worker per 256 concurrent rollouts). Only used when the orchestrator spawns the env server (i.e. `address` is not set). (2026-03-25)
 - **`[model.vlm]` (NEW — replaces auto-detection)**: VLM mode is now opt-in via a `[model.vlm]` sub-config with required `vision_encoder_attr` and `language_model_attr` fields. There is no auto-detection — if you train a VLM, you must add `[model.vlm]`. Existing multimodal configs need the new section. See `docs/multimodal.md` for the table of known model attrs. (2026-03-24)
 - **`model.optimization_dtype` / `model.reduce_dtype` (VLM models, RL only)**: VLM dtype validation now only applies to RL training (`TrainerConfig`), not SFT. VLM models used with `sft` no longer require `optimization_dtype='bfloat16'` / `reduce_dtype='bfloat16'`. RL training still enforces both to match vLLM inference. (2026-03-24)
 - **`model.optimization_dtype` / `model.reduce_dtype` (VLM models)**: Added validation that VLM models must use `optimization_dtype='bfloat16'` and `reduce_dtype='bfloat16'` to match vLLM inference. Previously valid configs with `float32` (the default) are now rejected for VLM model names. Set both fields to `"bfloat16"` when training VLMs. (2026-03-21)

pyproject.toml

@@ -76,7 +76,7 @@ disagg = [
     "deep-gemm ; platform_machine == 'x86_64'",
     "nixl",
     "nixl-cu12 ; platform_machine == 'x86_64'",
-    "vllm-router",
+    "vllm-router ; platform_machine == 'x86_64'",
 ]
 
 [dependency-groups]
@@ -110,7 +110,7 @@ override-dependencies = [
 
 [tool.uv.sources]
 torch = { index = "pytorch-cu128" }
-verifiers = { git = "https://github.com/PrimeIntellect-ai/verifiers.git", rev = "f5580acd" }
+verifiers = { git = "https://github.com/PrimeIntellect-ai/verifiers.git", rev = "0db45e6" }
 torchtitan = { git = "https://github.com/pytorch/torchtitan", rev = "a1fdd7e" }
 dion = { git = "https://github.com/samsja/dion.git", rev = "d891eeb" }
 transformers = { git = "https://github.com/huggingface/transformers.git", rev = "5c1c72b" }

scripts/slurm_tmux.sh

@@ -21,24 +21,30 @@ echo "Creating tmux session: $SESSION_NAME"
 # Window 0: Terminal
 tmux new-session -d -s "$SESSION_NAME" -n "Terminal"
 
-# Window 1: Logs - 3 vertical panes
+# Window 1: Logs - 4 vertical panes
 tmux new-window -t "$SESSION_NAME" -n "Logs"
 
 tmux split-window -v -t "$SESSION_NAME:Logs.0"
 tmux split-window -v -t "$SESSION_NAME:Logs.1"
+tmux split-window -v -t "$SESSION_NAME:Logs.2"
 tmux select-layout -t "$SESSION_NAME:Logs" even-vertical
 
 tmux select-pane -t "$SESSION_NAME:Logs.0" -T "Trainer"
 tmux select-pane -t "$SESSION_NAME:Logs.1" -T "Orchestrator"
-tmux select-pane -t "$SESSION_NAME:Logs.2" -T "Inference"
+tmux select-pane -t "$SESSION_NAME:Logs.2" -T "Envs"
+tmux select-pane -t "$SESSION_NAME:Logs.3" -T "Inference"
 
 tmux send-keys -t "$SESSION_NAME:Logs.0" \
   "tail -F ${SLURM_LOG_DIR}/latest_train_node_rank_*.log 2>/dev/null" C-m
 
 tmux send-keys -t "$SESSION_NAME:Logs.1" \
   "tail -F ${SLURM_LOG_DIR}/latest_orchestrator.log 2>/dev/null" C-m
 
+ENV_LOG_DIR="${OUTPUT_DIR}/logs/envs"
 tmux send-keys -t "$SESSION_NAME:Logs.2" \
+  "tail -F ${ENV_LOG_DIR}/*/*/*.log 2>/dev/null" C-m
+
+tmux send-keys -t "$SESSION_NAME:Logs.3" \
   "tail -F ${SLURM_LOG_DIR}/latest_infer_node_rank_*.log 2>/dev/null" C-m
 
 # Pane title styling

scripts/tmux.sh

@@ -58,23 +58,30 @@ else
   # Start new tmux session with first window
   tmux new-session -d -s "$SESSION_NAME" -n "RL"
 
-  # Window 1: RL - 3 vertical panes
+  # Window 1: RL - 4 vertical panes
   tmux split-window -v -t "$SESSION_NAME:RL.0"
   tmux split-window -v -t "$SESSION_NAME:RL.1"
+  tmux split-window -v -t "$SESSION_NAME:RL.2"
   tmux select-layout -t "$SESSION_NAME:RL" even-vertical
 
   # Pane titles
   tmux select-pane -t "$SESSION_NAME:RL.0" -T "Trainer"
   tmux select-pane -t "$SESSION_NAME:RL.1" -T "Orchestrator"
-  tmux select-pane -t "$SESSION_NAME:RL.2" -T "Inference"
+  tmux select-pane -t "$SESSION_NAME:RL.2" -T "Envs"
+  tmux select-pane -t "$SESSION_NAME:RL.3" -T "Inference"
 
   # Logs: Orchestrator
   tmux send-keys -t "$SESSION_NAME:RL.1" \
     "echo \"Following orchestrator.stdout (tail -F; waits for rotate/create)...\"; tail -F \"${OUTPUT_DIR}/logs/orchestrator.stdout\" 2>/dev/null" \
     C-m
 
-  # Logs: Inference
+  # Logs: Envs (all env server and worker logs)
   tmux send-keys -t "$SESSION_NAME:RL.2" \
+    "echo \"Following env logs (tail -F; waits for rotate/create)...\"; tail -F \"${OUTPUT_DIR}/logs/envs\"/*/*/*.log 2>/dev/null" \
+    C-m
+
+  # Logs: Inference
+  tmux send-keys -t "$SESSION_NAME:RL.3" \
     "echo \"Following inference.stdout (tail -F; waits for rotate/create)...\"; tail -F \"${OUTPUT_DIR}/logs/inference.stdout\" 2>/dev/null" \
     C-m
 

src/prime_rl/configs/env_server.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Annotated
 
-from pydantic import Field
+from pydantic import Field, model_validator
 
 from prime_rl.configs.orchestrator import EnvConfig
 from prime_rl.configs.shared import LogConfig
@@ -19,4 +19,10 @@ class EnvServerConfig(BaseConfig):
         Field(
             description="Directory to write outputs to. Will be populated with checkpoints, weights, rollouts and logs as subdirectories. Should be set to a persistent directory with enough disk space. This value should be distinct across experiments running on a single node. See the README for more details."
         ),
-    ] = Path("outputs/run_default")
+    ] = Path("outputs")
+
+    @model_validator(mode="after")
+    def validate_num_workers(self):
+        if self.env.num_workers == "auto":
+            self.env.num_workers = 1
+        return self

src/prime_rl/configs/orchestrator.py

@@ -287,6 +287,17 @@ class EnvConfig(BaseConfig):
             ),
         ),
     ] = {}
+    num_workers: Annotated[
+        int | Literal["auto"],
+        Field(
+            description=(
+                "Number of env server worker processes. "
+                "Set to 'auto' to scale based on the env's concurrency (1 worker per 256 concurrent rollouts). "
+                "When setting manually, we recommend sizing so that each worker handles at most 256 concurrent rollouts. "
+                "Only used when the orchestrator spawns the env server (i.e. address is None)."
+            ),
+        ),
+    ] = "auto"
     max_retries: Annotated[
         int,
         Field(

src/prime_rl/entrypoints/rl.py

@@ -455,6 +455,38 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) ->
     script_path.write_text(script)
 
 
+def format_log_message(
+    trainer_log: str,
+    orchestrator_log: str | None,
+    inference_log: str | None,
+    env_log_dir: Path,
+    train_env_names: list[str],
+    eval_env_names: list[str],
+) -> str:
+    col = 18
+    i1 = " " * 2
+    i2 = " " * 3
+    i3 = " " * 4
+    max_name = col - 4
+
+    log_lines = [f"{i1}{'Trainer:':<{col}}tail -F {trainer_log}"]
+    if orchestrator_log:
+        log_lines.append(f"{i1}{'Orchestrator:':<{col}}tail -F {orchestrator_log}")
+    if inference_log:
+        log_lines.append(f"{i1}{'Inference:':<{col}}tail -F {inference_log}")
+    log_lines.append(f"{i1}{'Envs:':<{col}}tail -F {env_log_dir}/*/*/*.log")
+    log_lines.append(f"{i2}{'Train:':<{col - 1}}tail -F {env_log_dir}/train/*/*.log")
+    for name in train_env_names:
+        short = name if len(name) <= max_name else name[: max_name - 3] + "..."
+        log_lines.append(f"{i3}{f'{short}:':<{col - 2}}tail -F {env_log_dir}/train/{name}/*.log")
+    if eval_env_names:
+        log_lines.append(f"{i2}{'Eval:':<{col - 1}}tail -F {env_log_dir}/eval/*/*.log")
+        for name in eval_env_names:
+            short = name if len(name) <= max_name else name[: max_name - 3] + "..."
+            log_lines.append(f"{i3}{f'{short}:':<{col - 2}}tail -F {env_log_dir}/eval/{name}/*.log")
+    return "Logs:\n" + "\n".join(log_lines)
+
+
 def rl_slurm(config: RLConfig):
     assert config.slurm is not None
 
@@ -466,22 +498,36 @@ def rl_slurm(config: RLConfig):
         logger.info(f"Wrote config to {config_dir / RL_TOML}")
 
         log_dir = get_log_dir(config.output_dir)
-        log_message = (
-            f"Logs:\n"
-            f"  Trainer:       tail -F {log_dir}/trainer.stdout\n"
-            f"  Orchestrator:  tail -F {log_dir}/orchestrator.stdout\n"
-            f"  Inference:     tail -F {log_dir}/inference.stdout"
+        env_log_dir = get_log_dir(config.output_dir) / "envs"
+        train_env_names = [env.resolved_name for env in config.orchestrator.env]
+        eval_env_names = [env.resolved_name for env in config.orchestrator.eval.env] if config.orchestrator.eval else []
+
+        log_message = format_log_message(
+            trainer_log=f"{log_dir}/trainer.stdout",
+            orchestrator_log=f"{log_dir}/orchestrator.stdout",
+            inference_log=f"{log_dir}/inference.stdout",
+            env_log_dir=env_log_dir,
+            train_env_names=train_env_names,
+            eval_env_names=eval_env_names,
         )
     else:
         write_subconfigs(config, config_dir)
         logger.info(f"Wrote subconfigs to {config_dir}")
 
         slurm_log_dir = config.output_dir / "slurm"
-        log_lines = [f"  Trainer:       tail -F {slurm_log_dir}/latest_train_node_rank_0.log"]
-        if config.deployment.num_infer_nodes > 0:
-            log_lines.append(f"  Orchestrator:  tail -F {slurm_log_dir}/latest_orchestrator.log")
-            log_lines.append(f"  Inference:     tail -F {slurm_log_dir}/latest_infer_node_rank_0.log")
-        log_message = "Logs:\n" + "\n".join(log_lines)
+        env_log_dir = get_log_dir(config.output_dir) / "envs"
+        train_env_names = [env.resolved_name for env in config.orchestrator.env]
+        eval_env_names = [env.resolved_name for env in config.orchestrator.eval.env] if config.orchestrator.eval else []
+
+        has_infer = config.deployment.num_infer_nodes > 0
+        log_message = format_log_message(
+            trainer_log=f"{slurm_log_dir}/latest_train_node_rank_0.log",
+            orchestrator_log=f"{slurm_log_dir}/latest_orchestrator.log" if has_infer else None,
+            inference_log=f"{slurm_log_dir}/latest_infer_node_rank_0.log" if has_infer else None,
+            env_log_dir=env_log_dir,
+            train_env_names=train_env_names,
+            eval_env_names=eval_env_names,
+        )
 
     script_path = config.output_dir / RL_SBATCH
     write_slurm_script(config, config_dir, script_path)

src/prime_rl/orchestrator/env_server/env_server.py

@@ -1,6 +1,6 @@
 import asyncio
 
-from verifiers.workers import ZMQEnvServer
+from verifiers.serve import ZMQEnvServer
 
 from prime_rl.configs.env_server import EnvServerConfig
 from prime_rl.utils.config import cli
@@ -19,17 +19,16 @@ def run_server(config: EnvServerConfig):
     for env_id in env_ids_to_install:
         install_env(env_id)
 
-    env_name = config.env.name or config.env.id
-    log_file = (get_log_dir(config.output_dir) / "train" / f"{env_name}.log").as_posix()
+    log_dir = (get_log_dir(config.output_dir) / config.env.resolved_name).as_posix()
 
     server = ZMQEnvServer(
         env_id=strip_env_version(config.env.id),
         env_args=config.env.args,
         extra_env_kwargs=config.env.extra_env_kwargs,
         log_level=config.log.level,
-        log_file_level=config.log.vf_level,
-        log_file=log_file,
+        log_dir=log_dir,
         json_logging=config.log.json_logging,
+        num_workers=config.env.num_workers,
         **{"address": config.env.address} if config.env.address is not None else {},
     )
     asyncio.run(server.run())

src/prime_rl/orchestrator/orchestrator.py

@@ -53,6 +53,7 @@
     get_completion_len,
     get_seq_len,
     intercept_vf_logging,
+    resolve_num_workers,
     setup_env_client,
     spawn_env_server,
     task_uses_group_scoring,
@@ -85,7 +86,7 @@ async def orchestrate(config: OrchestratorConfig):
         log_file=config.output_dir / "logs" / "orchestrator.log" if config.log.file else None,
         json_logging=config.log.json_logging,
     )
-    intercept_vf_logging(logger="verifiers.workers", level=config.log.vf_level)  # show logs from env clients
+    intercept_vf_logging(logger="verifiers.serve", level=config.log.vf_level)  # show logs from env clients
     logger.info("Starting orchestrator")
 
     event_loop_lag_monitor = EventLoopLagMonitor()
@@ -172,10 +173,10 @@ async def orchestrate(config: OrchestratorConfig):
 
     # Load environment and extract dataset
     logger.info(
-        f"Loading {len(config.env)} training environment(s) ({', '.join(env.name or env.id for env in config.env)})"
+        f"Loading {len(config.env)} training environment(s) ({', '.join(env.resolved_name for env in config.env)})"
     )
     env_ids = [strip_env_version(env.id) for env in config.env]
-    train_env_names = [env.name or env_id for env_id, env in zip(env_ids, config.env)]
+    train_env_names = [env.resolved_name for env in config.env]
     train_env_group = vf.EnvGroup(
         envs=[vf.load_environment(env_id, **env.args) for env_id, env in zip(env_ids, config.env)],
         env_names=train_env_names,
@@ -216,15 +217,18 @@ def _cleanup_env_processes():
 
     for env_id, env, env_name in zip(env_ids, config.env, train_env_names):
         if env.address is None:
+            num_workers = resolve_num_workers(env.num_workers, config.max_inflight_rollouts)
+            log_dir = (get_log_dir(config.output_dir.parent) / "envs" / "train" / env_name).as_posix()
             address, process = spawn_env_server(
                 env_id=env_id,
                 env_args=env.args,
                 extra_env_kwargs=env.extra_env_kwargs,
-                log_level="CRITICAL",
-                log_file=(get_log_dir(config.output_dir) / "train" / f"{env_name}.log").as_posix(),
-                log_file_level=config.log.vf_level,
+                num_workers=num_workers,
+                log_level=config.log.vf_level,
+                log_dir=log_dir,
                 json_logging=config.log.json_logging,
             )
+            logger.info(f"Spawned env server for {env_name} with {num_workers} worker(s)")
             env_processes.append(process)
         else:
             if env_name in train_env_deferred_group_scoring_tasks:
@@ -251,21 +255,34 @@ def _cleanup_env_processes():
     if config.eval:
         env_ids = [strip_env_version(env.id) for env in config.eval.env]
         eval_envs = [vf.load_environment(env_id, **env.args) for env_id, env in zip(env_ids, config.eval.env)]
-        eval_env_names = [env.name or env_id for env_id, env in zip(env_ids, config.eval.env)]
+        eval_env_names = [env.resolved_name for env in config.eval.env]
         eval_sampling_args = get_eval_sampling_args(config.eval.sampling)
         eval_env_addresses = []
 
         for env_id, env, eval_env_name in zip(env_ids, config.eval.env, eval_env_names):
             if env.address is None:
+                num_examples = env.num_examples or config.eval.num_examples
+                rollouts_per_example = env.rollouts_per_example or config.eval.rollouts_per_example
+                if num_examples == -1:
+                    max_concurrent = 1024
+                    logger.warning(
+                        f"Eval env '{eval_env_name}' uses all examples (num_examples=-1). "
+                        f"Defaulting max_concurrent={max_concurrent} for worker scaling."
+                    )
+                else:
+                    max_concurrent = num_examples * rollouts_per_example
+                num_workers = resolve_num_workers(env.num_workers, max_concurrent)
+                log_dir = (get_log_dir(config.output_dir.parent) / "envs" / "eval" / eval_env_name).as_posix()
                 address, process = spawn_env_server(
                     env_id=env_id,
                     env_args=env.args,
                     extra_env_kwargs=env.extra_env_kwargs,
-                    log_level="CRITICAL",
-                    log_file=(get_log_dir(config.output_dir) / "eval" / f"{eval_env_name}.log").as_posix(),
-                    log_file_level=config.log.vf_level,
+                    num_workers=num_workers,
+                    log_level=config.log.vf_level,
+                    log_dir=log_dir,
                     json_logging=config.log.json_logging,
                 )
+                logger.info(f"Spawned eval env server for {eval_env_name} with {num_workers} worker(s)")
                 env_processes.append(process)
             else:
                 address = env.address
@@ -490,9 +507,8 @@ def _cleanup_env_processes():
 
         # Schedule generating the training batch
         temperature = compute_temperature(progress.step, config.sampling, config.max_steps)
-        sampling_args = get_sampling_args(
-            config.sampling, temperature=temperature, use_token_client=config.use_token_client
-        )
+        is_vllm = config.teacher_rollout_model is None
+        sampling_args = get_sampling_args(config.sampling, temperature=temperature, is_vllm=is_vllm)
         scheduler.set_sampling_args(sampling_args)
         train_task = asyncio.create_task(scheduler.generate_batch(step=progress.step))
 

src/prime_rl/orchestrator/scheduler.py

@@ -88,9 +88,8 @@ def __init__(
         self.enable_policy_updates = enable_policy_updates
         self.lora_name = lora_name
         initial_temp = compute_temperature(step=0, sampling_config=config.sampling, max_steps=config.max_steps)
-        self.sampling_args = get_sampling_args(
-            config.sampling, temperature=initial_temp, use_token_client=config.use_token_client
-        )
+        is_vllm = config.teacher_rollout_model is None
+        self.sampling_args = get_sampling_args(config.sampling, temperature=initial_temp, is_vllm=is_vllm)
         self.model_name = self.config.model.name
         self.json_logging = config.log.json_logging
 

src/prime_rl/orchestrator/utils.py

@@ -36,13 +36,14 @@ async def get_semaphore() -> AsyncContextManager:
     return SEMAPHORE
 
 
-def get_sampling_args(sampling_config: SamplingConfig, temperature: float, use_token_client: bool = True) -> dict:
+def get_sampling_args(sampling_config: SamplingConfig, temperature: float, is_vllm: bool = True) -> dict:
     # Convert SamplingConfig to vLLM OAI sampling args
     # https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters_2
     sampling_args = dict(sampling_config)
     sampling_args.pop("temp_scheduler", None)
     sampling_args["temperature"] = temperature
     sampling_args["top_p"] = 1.0
+    sampling_args["logprobs"] = True
     extra_body = dict(sampling_config.extra_body)
 
     min_tokens = sampling_args.pop("min_tokens")
@@ -53,12 +54,9 @@ def get_sampling_args(sampling_config: SamplingConfig, temperature: float, use_t
     if repetition_penalty != 1.0:
         extra_body["repetition_penalty"] = repetition_penalty
 
-    extra_body["top_k"] = -1
-    extra_body["min_p"] = 0.0
-
-    sampling_args["logprobs"] = True
-
-    if use_token_client:
+    if is_vllm:
+        extra_body["top_k"] = -1
+        extra_body["min_p"] = 0.0
         extra_body["return_token_ids"] = True
 
     if extra_body: