code/train/ctrl_policy.py

"""Controller policy of the decoupled plan (PLAN_DECOUPLED_20261013.md sections 3 and 5): a frozen perception gives
(p, h, v) of the current target, a small MLP on the last 20 steps picks the MineRL action. Runs in the eval venv
(numpy, openai; no torch). MineExplorer/eval_benchmark.py (MC_POLICY) and the GRPO rollouts call Policy.act once per
step; train/ctrl.py (torch) trains the MLP and exports it as npz: w0, b0, w1, b1, w2, b2 with weights as (in, out).

Actions, three heads sampled independently:
- movement, 15 combinations (MOVES); sprint only with forward, else MineExplorer's validation swaps in a no-op. A VPT
  key bitmask maps to its class by cancelling opposite keys (as the game does), then dropping sneak, sprint and jump
  in turn until the combination is one of MOVES;
- camera yaw, 17 bins (YAWS: 0, +-0.5 ... +-64 deg, > 0 turns right); camera pitch, 13 bins (PITCHES: 0, +-0.5 ...
  +-16 deg, > 0 looks down). A recorded delta maps to the nearest bin on a log scale, 0 only for an exact 0.
attack, use, inventory, hotbar, ESC and the remaining keys are always 0.
Input: 5 numbers per step, the last 20 steps oldest first (100): p, h/HALF_HFOV, v/HALF_VFOV of the current target
(bench/diag_perception.py view: h > 0 right of the view center, v > 0 above it) and the policy's own previous camera
action yaw/64, pitch/16. Steps before the episode start are 0; when the target changes, the target part of the
earlier steps is set to 0. A controller whose w0 has 116 rows also gets 16 memory inputs after the window (memory():
for the 8 directions yaw + 45 i, i = 0..7, whether the direction's world-yaw sector was a view sector and whether the
4x4 cell 4 blocks along it was stood in, over the episode's poses up to the current one).
Perception:
- gt: p = 1 and the exact h, v when the target is in the geometric field of view (|h| <= 51.2, |v| <= 35), else 0;
- vlm: two questions, sent at the same time, on the frame the eval client sends (320x180 PNG), thinking off,
  temperature 0, worded as measured: step 1's grounding question (bench/diag_perception_direction.py GROUND,
  max_tokens 48) with the answer started as START, whose first box's center through the pinhole model gives h, v
  (0 without a box), and step 1b's presence score s_yes (bench/diag_perception_presence.py YESNO, max_tokens 1:
  P(yes) / (P(yes) + P(no)) from the top 20 logprobs) as p. Asked on the episode's first step, when the target name
  changes and PERCEIVE_EVERY steps after the last call; the steps in between carry the last answer over (carry()).
Current target: the first milestone in order that has rules, was not satisfied at spawn and is not completed (the
checker marks it); its name is the milestone task "find X" -> X with underscores as spaces.
Per-step log (Policy.act): x (the input, rounded to 4 decimals; the forward pass uses exactly these values), target
(milestone index), bins [move, yaw, pitch], logp (sum over the three heads), perception (vlm: reply, box, p on the
steps that asked, None on the carried ones; gt: None), held (steps since the last VLM call: 0 on the steps that asked,
and always for gt), truth of the current target (seen = in the geometric field of view, dist from the feet as the
milestone check, h, v; the checker's yaw error is -h).
"""
from collections import deque
from concurrent.futures import ThreadPoolExecutor
import math
import time
from pathlib import Path
import sys

import numpy as np
from openai import OpenAI

ROOT = Path(__file__).resolve().parents[1]
sys.path.append(str(ROOT / "bench"))
from diag_perception import EYE, HALF_VFOV, view  # noqa: E402
from diag_perception_direction import BOX, GROUND, TAN_HALF_HFOV  # noqa: E402
from diag_perception_presence import YESNO  # noqa: E402

WINDOW = 20
START = '```json\n[\n\t{"bbox_2d": ['  # the start of the grounding answer (ground())
PERCEIVE_EVERY = 3  # vlm: steps from one perception call to the next (find25 evals and GRPO rollouts)
HALF_HFOV = math.degrees(math.atan(TAN_HALF_HFOV))  # 51.2
MOVES = [(), ("forward",), ("forward", "jump"), ("forward", "right"), ("forward", "left"), ("forward", "sprint"),
         ("forward", "jump", "sprint"), ("back",), ("right",), ("left",), ("back", "left"), ("back", "right"),
         ("forward", "right", "jump"), ("jump",), ("forward", "left", "jump")]
YAWS = np.array([-64, -32, -16, -8, -4, -2, -1, -0.5, 0, 0.5, 1, 2, 4, 8, 16, 32, 64])
PITCHES = np.array([-16, -8, -4, -2, -1, -0.5, 0, 0.5, 1, 2, 4, 8, 16])
CAM_SCALE = np.array([16.0, 64.0])  # camera [pitch, yaw] -> input
KEYS = {"forward": 1, "back": 2, "left": 4, "right": 8, "jump": 16, "sprint": 32, "sneak": 64}  # VPT key bitmask
NOOP = {"ESC": 0, "attack": 0, "back": 0, "forward": 0, "left": 0, "right": 0, "jump": 0, "camera": [0.0, 0.0],
        "drop": 0, **{f"hotbar.{i}": 0 for i in range(1, 10)}, "inventory": 0, "pickItem": 0, "sneak": 0,
        "sprint": 0, "swapHands": 0, "use": 0}


def move_class(keys):
    """MOVES index of a VPT key bitmask."""
    masks = [sum(KEYS[k] for k in m) for m in MOVES]
    for pair in (3, 12):  # forward+back, left+right
        if keys & pair == pair:
            keys &= ~pair
    for drop in (0, 64, 32, 16):
        keys &= ~drop
        if keys in masks:
            return masks.index(keys)


MOVE_OF_KEYS = np.array([move_class(k) for k in range(128)])


def cam_bin(deg, bins):
    """Bin index (into YAWS or PITCHES) of camera deltas in degrees."""
    j = np.clip(np.round(np.log2(np.maximum(np.abs(deg), 1e-9) / 0.5)), 0, len(bins) // 2 - 1)
    return (len(bins) // 2 + np.sign(deg) * (j + 1)).astype(int)


def gt_features(h, v):
    """GT perception (..., 3): p, h/HALF_HFOV, v/HALF_VFOV in the geometric field of view, else 0."""
    seen = (np.abs(h) <= HALF_HFOV) & (np.abs(v) <= HALF_VFOV)
    return np.stack([seen, h / HALF_HFOV * seen, v / HALF_VFOV * seen], -1).astype(float)


def rows(target, prev_cam):
    """Window steps (..., 5) from target features (..., 3) and the previous camera action [pitch, yaw] (..., 2)."""
    return np.concatenate([target, (prev_cam / CAM_SCALE)[..., ::-1]], -1)


def sector(yaw):
    """World-yaw sector of a yaw in degrees as logged (unwrapped too): 8 sectors of 45 deg, sector 0 = -22.5..22.5."""
    return int(((yaw + 22.5) % 360) // 45)


def cell(x, z):
    """4x4 XZ cell of a position, as bench/find_metrics.py."""
    return math.floor(x / 4), math.floor(z / 4)


def memory(sectors, cells, pose):
    """The 16 memory inputs at pose: for the directions t = yaw + 45 i, i = 0..7 (> 0 turns right), 1 if sector(t) is
    in sectors, then 1 if the cell 4 blocks along t (horizontal direction dx = -sin t, dz = cos t) is in cells."""
    yaws = [pose["yaw"] + 45 * i for i in range(8)]
    return np.array([sector(t) in sectors for t in yaws]
                    + [cell(pose["x"] - 4 * math.sin(math.radians(t)), pose["z"] + 4 * math.cos(math.radians(t)))
                       in cells for t in yaws], dtype=float)


def logits(w, x):
    """(move, yaw, pitch) logits for inputs x (..., 100, or 116 with memory); w = the exported npz."""
    z = np.maximum(x @ w["w0"] + w["b0"], 0)
    z = np.maximum(z @ w["w1"] + w["b1"], 0)
    return np.split(z @ w["w2"] + w["b2"], [len(MOVES), len(MOVES) + len(YAWS)], -1)


def log_softmax(z):
    z = z - z.max(-1, keepdims=True)
    return z - np.log(np.exp(z).sum(-1, keepdims=True))


def action(move, yaw, pitch):
    """MineRL action dict of the three bins."""
    a = dict(NOOP, camera=[float(PITCHES[pitch]), float(YAWS[yaw])])
    a.update({key: 1 for key in MOVES[move]})
    return a


def target_name(task):
    """Name asked about a milestone task: "find X" -> X with underscores as spaces (diagnoses B, C and step 1)."""
    return task[5:].removeprefix("a ").replace("_", " ")


def current_target(checker, presatisfied):
    """(milestone index, absolute target [x, y, z], name) of the current target."""
    for i, m in enumerate(checker._milestones):
        if m.get("rules") and m["milestone_id"] not in presatisfied and not checker._completed[i]:
            params = m["rules"][0]["params"]
            target = params["target"]
            if params.get("coordinate_frame") == "spawn_relative":
                o = checker._spawn_pos
                target = [target[0] + o["x"], target[1] + o["y"], target[2] + o["z"]]
            return i, target, target_name(m["task"])


def box_angles(box):
    """(h, v) of a box [x1, y1, x2, y2] in Qwen's relative 0-1000 coordinates: its center through the pinhole model
    (exact at pitch 0), as step 1."""
    x1, y1, x2, y2 = box
    return (math.degrees(math.atan(((x1 + x2) / 1000 - 1) * TAN_HALF_HFOV)),
            math.degrees(math.atan((1 - (y1 + y2) / 1000) * math.tan(math.radians(HALF_VFOV)))))


def chat(client, model, frame, prompt, max_tokens, start=None, **kwargs):
    """One question about one frame (base64 PNG), thinking off, temperature 0, as steps 1 and 1b asked it. start: the
    beginning of the answer, which the model continues (the reply is the continuation, its stop string included)."""
    messages = [{"role": "user", "content": [
        {"type": "image_url", "image_url": {"url": "data:image/png;base64," + frame}},
        {"type": "text", "text": prompt}]}]
    extra = {"chat_template_kwargs": {"enable_thinking": False}}
    if start:
        messages.append({"role": "assistant", "content": start})
        extra.update(continue_final_message=True, add_generation_prompt=False, include_stop_str_in_output=True)
    return client.chat.completions.create(model=model, temperature=0, max_tokens=max_tokens, messages=messages,
                                          extra_body=extra, **kwargs).choices[0]


def ground(client, model, frame, name):
    """Step 1's grounding question, max_tokens 48 -> (reply, first box or None, h, v); h = v = 0 without a box.
    The answer is started as START, how 15,505 of the first 15,533 replies of the BC find25 began (the other 28, 20
    of them about melon, said there is none and now get a box), so the model writes the first box's four numbers and
    stops at the first "]"; the prompt is then the question's followed by the tokens the model wrote for START
    (vLLM /tokenize). reply: START and the model's numbers."""
    reply = START + chat(client, model, frame, GROUND.format(name), 48, start=START, stop=["]"]).message.content
    box = BOX.search(reply)
    if not box:
        return reply, None, 0.0, 0.0
    box = [float(c) for c in box.groups()]
    return reply, box, *box_angles(box)


def presence(client, model, frame, name):
    """Step 1b's s_yes: diagnosis B's yes/no question, max_tokens 1, top 20 logprobs; P(yes) / (P(yes) + P(no)) with
    case and leading-space variants summed."""
    top = chat(client, model, frame, YESNO.format(name), 1, logprobs=True, top_logprobs=20).logprobs.content[0]
    p = {w: sum(math.exp(t.logprob) for t in top.top_logprobs if t.token.strip().lower() == w) for w in ("yes", "no")}
    return p["yes"] / (p["yes"] + p["no"])


def vlm_features(client, model, frame, name):
    """VLM perception -> (p, h/HALF_HFOV, v/HALF_VFOV), log {reply, box, p}; the two questions are sent at once."""
    with ThreadPoolExecutor(1) as pool:
        grounding = pool.submit(ground, client, model, frame, name)
        p = presence(client, model, frame, name)
        reply, box, h, v = grounding.result()
    return np.array([p, h / HALF_HFOV, v / HALF_VFOV]), {"reply": reply, "box": box, "p": round(p, 4)}


def carry(p, box, yaw, pitch, pose):
    """Features at pose of a VLM answer (p, first box or None) given at view yaw, pitch, as the C_vlm simulation keeps
    a box's world point (train/ctrl.py simulate()): h moves by -delta yaw, v by +delta pitch, the box center stays in
    the image; p is held, and without a box h = v = 0."""
    if box is None:
        return np.array([p, 0.0, 0.0])
    h, v = box_angles(box)
    h = np.clip(h - ((pose["yaw"] - yaw + 180) % 360 - 180), -HALF_HFOV, HALF_HFOV)
    v = np.clip(v + pose["pitch"] - pitch, -HALF_VFOV, HALF_VFOV)
    return np.array([p, h / HALF_HFOV, v / HALF_VFOV])


class Policy:
    """The controller for one episode. seed: anything np.random.default_rng takes (eval: [seed, task id]);
    perception "gt" or "vlm" (url and model of the perception server)."""

    def __init__(self, npz, seed, perception="gt", url=None, model=None):
        self.w = dict(np.load(npz))
        self.rng = np.random.default_rng(seed)
        self.perception, self.model = perception, model
        self.client = OpenAI(api_key="EMPTY", base_url=url) if perception == "vlm" else None
        self.window = deque((np.zeros(5) for _ in range(WINDOW)), maxlen=WINDOW)
        self.target, self.cam = None, [0.0, 0.0]
        self.asked, self.held = None, 0  # vlm: (name, p, box, yaw, pitch) of the last call; steps since it
        self.memory = len(self.w["w0"]) == 5 * WINDOW + 16
        self.sectors, self.cells = set(), set()  # memory: view sectors and cells of the episode's poses so far

    def vlm_features(self, frame, name):
        """The VLM perception call (train/optimus_policy.py ExpertPerceptionPolicy overrides it)."""
        return vlm_features(self.client, self.model, frame, name)

    def act(self, frame, info, checker, presatisfied):
        """frame: this step's frame as the eval client sends it (base64 PNG, 320x180; used by vlm); info: the env
        info before the action (player_pos); checker: the episode's MilestoneChecker; presatisfied: ids of milestones
        satisfied at spawn. Returns (MineRL action dict, step log)."""
        index, target, name = current_target(checker, presatisfied)
        if index != self.target:
            for r in self.window:
                r[:3] = 0
            self.target = index
        pose = info["player_pos"]
        dist, h, v = view(pose, target)
        if self.perception == "gt":
            feat, raw = gt_features(h, v), None
        elif self.asked and self.asked[0] == name and self.held < PERCEIVE_EVERY - 1:
            self.held += 1
            feat, raw = carry(*self.asked[1:], pose), None
        else:
            feat, raw = self.vlm_features(frame, name)
            self.asked, self.held = (name, feat[0], raw["box"], pose["yaw"], pose["pitch"]), 0
        self.window.append(rows(feat, np.array(self.cam)))
        x = np.concatenate(self.window)
        if self.memory:
            self.sectors.add(sector(pose["yaw"]))
            self.cells.add(cell(pose["x"], pose["z"]))
            x = np.concatenate([x, memory(self.sectors, self.cells, pose)])
        x = np.round(x, 4)
        heads = [log_softmax(z) for z in logits(self.w, x)]
        bins = [int(self.rng.choice(len(lp), p=np.exp(lp))) for lp in heads]
        a = action(*bins)
        self.cam = a["camera"]
        return a, {"x": x.tolist(), "target": index, "bins": bins,
                   "logp": float(sum(lp[b] for lp, b in zip(heads, bins))), "perception": raw, "held": self.held,
                   "truth": {"seen": bool(gt_features(h, v)[0]), "dist": round(dist, 2), "h": round(h, 2),
                             "v": round(v, 2)}}


# BT arm (EXPERIMENT_LOG.md 2026-10-17): the VLM (Qwen3.5-9B + LoRA) reads the current frame and the target name and
# writes a behavior code, one digit per FSQ dimension of train/bt_tokenizer.py; the next CHUNK ticks run open loop,
# each tick's three bins sampled from the decoder's distributions for the code (table.npz probs, mixed radix index).
CHUNK = 20
BT_PROMPT = "Task: find {}. Reply with the behavior code for the next second."
BT_MEMORY = " Seen directions (0, 45, ..., 315 deg right of view): {}. Visited cells 4 blocks along them: {}."
BT_NOOP = [0, len(YAWS) // 2, len(PITCHES) // 2]


def bt_prompt(name, mem=None):
    """The BT text prompt for target name; with mem (the two 8-character 0/1 strings of memory(), directions then
    cells) BT_MEMORY follows."""
    return BT_PROMPT.format(name) + (BT_MEMORY.format(*mem) if mem else "")


def bt_code(reply, levels):
    """Table index of a reply: its first len(levels) characters as digits, each clipped to its level; None if they
    are not all digits."""
    digits = reply[:len(levels)]
    if len(digits) < len(levels) or not digits.isdigit():
        return None
    index = 0
    for d, level in zip(digits, levels):
        index = index * level + min(int(d), level - 1)
    return index


class BTPolicy:
    """The BT policy for one episode, with Policy's act(). table: the tokenizer's decode-table npz; url, model: the
    vLLM serving the LoRA. seed: None (GRPO rollouts) or [eval seed, task id]; decision k of the episode then samples
    with seed eval_seed * 10**7 + task_id * 1000 + k, and the tick bins with a numpy Generator seeded by
    eval_seed * 10**7 + task_id * 1000. Asks on the episode's first step, every CHUNK ticks and when the current target
    changes; an unparseable reply runs CHUNK no-op ticks (valid False). memory True (the curiosity variant,
    EXPERIMENT_LOG.md 2026-10-18): every step adds the pose's sector and cell to the episode's sets as Policy does, and
    each decision's prompt carries memory() at that pose as two 0/1 strings (bt_prompt()).
    Per-step log: target, bins, tick (index in the chunk), truth as Policy; on decision steps also reply, code (the
    clipped digits), valid, logprobs (vLLM, per generated token), act_s (seconds of the call) and, with memory, memory
    (the two strings in the prompt)."""

    def __init__(self, table, seed, url, model, memory=False):
        t = np.load(table)
        self.levels, self.probs = [int(x) for x in t["levels"]], t["probs"]
        self.seed = None if seed is None else seed[0] * 10**7 + seed[1] * 1000
        self.rng = np.random.default_rng(self.seed)
        self.client, self.model = OpenAI(api_key="EMPTY", base_url=url), model
        self.target, self.tick, self.decisions, self.chunk = None, CHUNK, 0, None
        self.memory, self.sectors, self.cells = memory, set(), set()

    def ask(self, frame, name, mem=None):
        extra = {"chat_template_kwargs": {"enable_thinking": False}}
        if self.seed is not None:
            extra["seed"] = self.seed + self.decisions
        choice = self.client.chat.completions.create(
            model=self.model, temperature=1.0, max_tokens=len(self.levels), logprobs=True, extra_body=extra,
            messages=[{"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": "data:image/png;base64," + frame}},
                {"type": "text", "text": bt_prompt(name, mem)}]}]).choices[0]
        return choice.message.content, [t.logprob for t in choice.logprobs.content]

    def act(self, frame, info, checker, presatisfied):
        index, target, name = current_target(checker, presatisfied)
        pose = info["player_pos"]
        dist, h, v = view(pose, target)
        log = {"target": index}
        if self.memory:
            self.sectors.add(sector(pose["yaw"]))
            self.cells.add(cell(pose["x"], pose["z"]))
        if index != self.target or self.tick == CHUNK:
            t = time.monotonic()
            mem = None
            if self.memory:
                bits = "".join(str(int(b)) for b in memory(self.sectors, self.cells, pose))
                mem = log["memory"] = [bits[:8], bits[8:]]
            reply, logprobs = self.ask(frame, name, mem)
            code = bt_code(reply, self.levels)
            self.chunk = self.probs[code] if code is not None else None
            self.target, self.tick = index, 0
            self.decisions += 1
            log.update(reply=reply, valid=code is not None, logprobs=logprobs, act_s=round(time.monotonic() - t, 3),
                       code=code is not None and "".join(str(min(int(d), l - 1)) for d, l in zip(reply, self.levels))
                       or None)
        bins = BT_NOOP if self.chunk is None else [
            int(self.rng.choice(len(p), p=p / p.sum()))
            for p in np.split(self.chunk[self.tick].astype(float), [len(MOVES), len(MOVES) + len(YAWS)])]
        log.update(bins=bins, tick=self.tick, truth={"seen": bool(gt_features(h, v)[0]), "dist": round(dist, 2),
                                                     "h": round(h, 2), "v": round(v, 2)})
        self.tick += 1
        return action(*bins), log


class NoopPolicy:
    """Every step the no-op action (kind "noop")."""

    def act(self, frame, info, checker, presatisfied):
        return dict(NOOP), {}


def make_policy(kind="ctrl", **kwargs):
    """Policy (kind "ctrl") or BTPolicy (kind "bt") from an MC_POLICY / rollout job policy dict; kinds "hid" and
    "hid_teacher": the hidden-state arms' policies (train/hid_policy.py)."""
    if kind.startswith("hid"):
        from hid_policy import make
        return make(kind, **kwargs)
    if kind == "box":  # task 4b (EXPERIMENT_LOG.md 2026-10-19): train/box_policy.py
        from box_policy import BoxPolicy
        return BoxPolicy(**kwargs)
    if kind == "noop":  # what a 1B expert that never writes parseable JSON amounts to (EXPERIMENT_LOG.md 2026-10-19)
        return NoopPolicy()
    if kind == "openha":  # OpenHA baseline (EXPERIMENT_LOG.md 2026-10-19): train/openha_policy.py
        from openha_policy import OpenHAPolicy
        return OpenHAPolicy(**kwargs)
    if kind.startswith("optimus"):  # task 1A (EXPERIMENT_LOG.md 2026-10-19): train/optimus_policy.py
        import optimus_policy
        return (optimus_policy.OptimusActionPolicy if kind == "optimus_action"
                else optimus_policy.ExpertPerceptionPolicy)(**kwargs)
    return BTPolicy(**kwargs) if kind == "bt" else Policy(**kwargs)