Support branching rollouts via trajectories; refactor state handling (#549)

* big chungus refactor for branching rollouts + cleaner state handling

* tests passing

* 3.11 fix; ruff

* vllm logprob args

* dict indexing for messages

* remove generateinputs

* optional truncation in trajectorystep for tokens

* small tweaks

* optional decorator rank for sorting order

* minor tweak

* change rank -> priority

* add cleanup to is_completed

* tool_env error handling, sandbox command timeout

* handle updated context length msg

* duplicate is_truncated field

* add model/sampling to state

* client/model/sampling in init_state

* updated config

* add kimi overlong prompt message

* add kimi overlong prompt message

* set_max_seq_len

* Add numpy, sympy, and scipy to PythonEnv

* pin prime-rl to will/trajectories branch

* update prime-rl wiki-search config

* ruff, ty

* fix init_state tests

* version, release notes

* env version bumps

* ty fixes

* opt deps for ty CI

* pin trajectories for configs

* use verifiers commit for configs

* ty for verifiers only

* bump vllm version

* process overlong prompt into trajectory

* skip steps with None tokens

---------

Co-authored-by: Sebastian <[email protected]>

Author: willccbb

.github/workflows/style.yml

@@ -29,12 +29,12 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v6
         with:
-          python-version: '3.11'
+          python-version: "3.11"
       - name: Install uv
         uses: astral-sh/setup-uv@v4
         with:
           version: "latest"
       - name: Install dependencies
-        run: uv sync
+        run: uv sync --extra rl
       - name: Run ty
-        run: uv run ty check .
+        run: uv run ty check verifiers

.gitignore

@@ -29,7 +29,7 @@ docs/build/
 *.pyc
 
 # libraries
-prime-rl/
+prime-rl
 
 # outputs
 wandb/

configs/prime-rl/wiki-search.toml

@@ -1,13 +1,14 @@
-inference_gpu_ids = [0]
-trainer_gpu_ids = [1]
+inference_gpu_ids = [0,1,2,3,4,5]
+trainer_gpu_ids = [6,7]
 
 max_steps = 500
+max_async_level = 4
 
 [model]
-name = "Qwen/Qwen3-4B-Instruct-2507"
+name = "Qwen/Qwen3-4B-Thinking-2507"
 
 [wandb]
-project = "wiki-search"
+project = "wiki-search-debug"
 name = "wiki-search-4b"
 
 [trainer.optim]
@@ -31,16 +32,17 @@ target_modules = [
 [orchestrator]
 batch_size = 512
 rollouts_per_example = 16
-seq_len = 4096
+seq_len = 16384
 mask_truncated_completions = false
 zero_truncated_completions = true
+oversampling_factor = 2.0
+
 
 [orchestrator.sampling]
-max_tokens = 512
+max_tokens = 4096
 
 [orchestrator.buffer]
-type = "online-difficulty"
-oversampling_factor = 2.0
+online_difficulty_filtering = true
 
 [[orchestrator.env]]
 id = "primeintellect/wiki-search"

configs/vf-rl/reasoning-gym.toml

@@ -10,13 +10,11 @@ num_eval_examples = 2000
 seed = 1
 
 [inference]
-gpus = 4
-tensor_parallel_size = 2
-data_parallel_size = 2
+gpus = 6
 enforce_eager = true
 
 [trainer]
-gpus = 4
+gpus = 2
 batch_size = 512
 micro_batch_size = 2
 max_seq_len = 4096

configs/vf-rl/wiki-search.toml

@@ -1,7 +1,7 @@
 model = "Qwen/Qwen3-4B-Instruct-2507"
 
 [env]
-id = "primeintellect/wiki-search"
+id = "wiki-search"
 
 [env.args]
 max_turns = 10
@@ -20,7 +20,7 @@ gpus = 1
 run_name = "wiki-search"
 micro_batch_size = 4
 rollouts_per_example = 16
-batch_size = 1024
+batch_size = 512
 max_steps = 500
 max_tokens = 512
 max_seq_len = 4096

configs/vf-rl/wordle.toml

@@ -1,21 +1,18 @@
 model = "Qwen/Qwen3-4B-Instruct-2507"
 
 [env]
-id = "will/wordle"
+id = "wordle"
 
 [inference]
 gpus = 1
 
-[inference.args]
-enforce_eager = true
-
 [trainer]
 gpus = 1
 
 [trainer.args]
 lora_target_modules = "all-linear"
 run_name = "wordle"
-micro_batch_size = 8
+micro_batch_size = 4
 rollouts_per_example = 16
 batch_size = 512
 max_steps = 500

environments/math_group/pyproject.toml

@@ -1,8 +1,8 @@
 [project]
 name = "math-group"
-version = "0.1.0"
+version = "0.1.1"
 dependencies = [
-    "verifiers>=0.1.4",
+    "verifiers>=0.1.8",
     "math-verify>=0.8.0",
 ]
 

environments/math_python/pyproject.toml

@@ -2,10 +2,10 @@
 name = "math-python"
 description = "Solve math problems using Python in a sandbox environment"
 tags = ["tool-use", "math", "sandbox", "train", "prime-sandboxes", "python", "coding"]
-version = "0.1.7"
+version = "0.1.8"
 requires-python = ">=3.11"
 dependencies = [
-    "verifiers>=0.1.5.post0",
+    "verifiers>=0.1.8",
     "math-verify>=0.8.0",
 ]
 

environments/sentence_repeater/sentence_repeater.py

@@ -1,7 +1,7 @@
 import random
 from copy import deepcopy
 from difflib import SequenceMatcher
-from typing import List, Tuple
+from typing import List
 
 from datasets import Dataset, load_dataset
 
@@ -75,19 +75,19 @@ class SentenceRepeaterEnv(vf.MultiTurnEnv):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
-    async def is_completed(self, messages: Messages, state: State, **kwargs) -> bool:
-        max_turns_reached = await super().is_completed(messages, state, **kwargs)
-        return state["turn"] >= len(state["info"]["questions"]) or max_turns_reached
+    @vf.stop
+    async def all_questions_answered(self, state: State) -> bool:
+        return len(state["trajectory"]) >= len(state["info"]["questions"])
 
     async def env_response(
         self, messages: Messages, state: State, **kwargs
-    ) -> Tuple[Messages, State]:
+    ) -> Messages:
         return [
             {
                 "role": "user",
                 "content": state["info"]["questions"][state["turn"]],
             }
-        ], state
+        ]
 
 
 def load_environment(**kwargs) -> vf.Environment:

environments/wiki_search/pyproject.toml

@@ -3,9 +3,9 @@ name = "wiki-search"
 description = "Agentic RAG over Wikipedia pages for trivia Q&A"
 tags = ["wikipedia", "multi-turn", "agentic-search", "rag", "train", "eval", "llm-judge"]
 requires-python = ">=3.11"
-version = "0.1.20"
+version = "0.1.21"
 dependencies = [
-    "verifiers>=0.1.7",
+    "verifiers>=0.1.8",
     "chromadb",
     "datasets",
     "openai",