fix(rollout): replace generate_rollout_completions (vLLM-only) with direct model.generate()

generate_rollout_completions hard-crashes when use_vllm=False. Replace with
_generate_no_vllm() which uses HF model.generate() + output_scores=True to
get prompt_ids, completion_ids, logprobs, and text without vLLM.

training/rollout.py

@@ -16,11 +16,59 @@ import json
 import re
 from typing import Any
 
-from trl.experimental.openenv import generate_rollout_completions
+import torch
+import torch.nn.functional as F
 
 from training.dataset import parse_seed_from_prompt
 from training.prompts import SYSTEM_PROMPT, format_observation
 
+
+# ---------------------------------------------------------------------------
+# HF model.generate() — replaces generate_rollout_completions (vLLM-only)
+# ---------------------------------------------------------------------------
+
+def _generate_no_vllm(trainer, prompt_text: str, tokenizer, max_new_tokens: int = 512) -> dict:
+    """Generate one completion using HF model.generate() without vLLM.
+
+    Returns the same dict shape as generate_rollout_completions so the rest
+    of rollout_once is unchanged:
+        prompt_ids:     list[int]
+        completion_ids: list[int]
+        logprobs:       list[float]  (per-token log-prob under current policy)
+        text:           str
+    """
+    device = next(trainer.model.parameters()).device
+    enc = tokenizer(prompt_text, return_tensors="pt").to(device)
+    prompt_len = enc["input_ids"].shape[1]
+
+    with torch.no_grad():
+        out = trainer.model.generate(
+            **enc,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+            output_scores=True,
+            return_dict_in_generate=True,
+        )
+
+    prompt_ids     = enc["input_ids"][0].tolist()
+    completion_ids = out.sequences[0][prompt_len:].tolist()
+
+    # Per-token log-probs from output.scores (one score tensor per new token)
+    logprobs = [
+        F.log_softmax(score[0], dim=-1)[tok_id].item()
+        for score, tok_id in zip(out.scores, completion_ids)
+    ]
+
+    text = tokenizer.decode(completion_ids, skip_special_tokens=True)
+    return {
+        "prompt_ids":     prompt_ids,
+        "completion_ids": completion_ids,
+        "logprobs":       logprobs,
+        "text":           text,
+    }
+
 MAX_STEPS = 40
 # ~3000 tokens at 4 chars/token; leaves room for completion tokens
 MAX_PROMPT_CHARS = 12_000
@@ -209,14 +257,12 @@ def rollout_once(
             enable_thinking=False,
         )
 
-        rollout_out = generate_rollout_completions(trainer, [prompt_text])[0]
+        rollout_out = _generate_no_vllm(trainer, prompt_text, tokenizer)
         prompt_ids.extend(rollout_out["prompt_ids"])
         completion_ids.extend(rollout_out["completion_ids"])
         logprobs.extend(rollout_out["logprobs"])
 
-        completion_text = rollout_out.get("text") or tokenizer.decode(
-            rollout_out["completion_ids"], skip_special_tokens=True
-        )
+        completion_text = rollout_out["text"]
 
         # Parse action; fall back gracefully on parse failure
         parsed = extract_json_action(completion_text)