training_code/model_organisms/envs/base.py

"""Shared types for the RL environments.

Design notes
------------
* We build the token sequence + loss mask *incrementally* and exactly: assistant
  segments are the policy's own generated token ids (never re-tokenized text), and
  Qwen3 chat control tokens / observations are appended as untrained tokens. This
  avoids the classic multi-turn re-tokenization drift.
* Generation always stops *before* ``<|im_end|>`` (stop string), so we append the
  turn-closing control tokens ourselves with a uniform structure.
* ``enable_thinking=True``: we *want* a CoT (that is what gets obfuscated). The
  think length is bounded by the think-budget logits processor at generation time.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

IM_END = "<|im_end|>\n"


@dataclass
class RLTrajectory:
    """One sampled rollout: token sequence + which tokens the policy generated."""

    input_ids: list[int]
    assistant_mask: list[bool]  # True where token is policy-generated (trained on)
    task_reward: float          # R_task in [0,1]
    m_out: float                # output-monitor score in [0,1]
    m_cot: float                # CoT-monitor score in [0,1] (eval only, not in reward)
    prompt_text: str
    think_text: str
    output_text: str
    full_text: str
    group_id: int               # which prompt/group this belongs to
    meta: dict[str, Any] = field(default_factory=dict)
    # filled by the trainer:
    old_logprob: float = 0.0
    advantage: float = 0.0
    reward: float = 0.0


class SeqBuilder:
    """Accumulate token ids + trained-token mask across turns of one rollout."""

    def __init__(self, tokenizer, prefix_ids: list[int]):
        self.tok = tokenizer
        self.ids: list[int] = list(prefix_ids)
        self.mask: list[bool] = [False] * len(prefix_ids)

    def add_generated(self, gen_ids: list[int]) -> None:
        self.ids.extend(gen_ids)
        self.mask.extend([True] * len(gen_ids))

    def add_control(self, text: str) -> None:
        toks = self.tok.encode(text, add_special_tokens=False)
        self.ids.extend(toks)
        self.mask.extend([False] * len(toks))

    def close_assistant(self) -> None:
        """Append the turn-closing ``<|im_end|>`` (untrained)."""
        self.add_control(IM_END)

    def add_user_turn(self, content: str) -> None:
        """Close prior assistant turn already done; add a user turn + open next assistant."""
        self.add_control(f"<|im_start|>user\n{content}{IM_END}<|im_start|>assistant\n")


def initial_prefix_ids(tokenizer, user_content: str, system: str | None = None) -> list[int]:
    """Tokenized chat prefix ending in the assistant generation prompt."""
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": user_content})
    out = tokenizer.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True, enable_thinking=True, return_dict=False
    )
    if hasattr(out, "input_ids"):  # BatchEncoding fallback
        out = out.input_ids
    if out and isinstance(out[0], list):
        out = out[0]
    return list(out)


def split_think_output(text: str) -> tuple[str, str]:
    """Split assistant content into (think, output) around ``</think>``."""
    if "</think>" in text:
        think, _, output = text.partition("</think>")
        return think.replace("<think>", "").strip(), output.strip()
    return text.replace("<think>", "").strip(), ""