Fix: decode with skip_special_tokens=False to preserve <think> tags for stripping

training/rollout.py

@@ -121,8 +121,14 @@ def rollout_episode(
             pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
         )
         completion_ids = gen[0, prompt_ids.shape[0]:]
-        raw_text = tokenizer.decode(completion_ids, skip_special_tokens=True)
-        completion_text = strip_think_blocks(raw_text)
+        # Decode WITHOUT skipping special tokens so <think>...</think> tags
+        # are preserved for our regex. Then strip think blocks, then remove
+        # remaining special tokens (EOS, chat markers, etc.).
+        raw_text = tokenizer.decode(completion_ids, skip_special_tokens=False)
+        cleaned = strip_think_blocks(raw_text)
+        for tok in tokenizer.all_special_tokens:
+            cleaned = cleaned.replace(tok, "")
+        completion_text = cleaned.strip()
         del gen
         if torch.cuda.is_available():
             torch.cuda.empty_cache()