src/mcss/vision_probe/opportunity_selectors.py
"""Mask-free visible signals and discovery-only, locked opportunity selectors.
Matching confidence is computed over every token, including background. The
selection RGB partition is observed data, not independent evaluation data.
"""
from __future__ import annotations
import numpy as np
import torch
from torch import Tensor
from torch.nn import functional as F
from mcss.vision_probe.adaptation import ACTIONS, FastState, Readout, materialize
SCHEMA = "mcss.opportunity_selectors.v1"
ALPHAS = (0.1, 1.0, 10.0, 100.0)
THRESHOLDS = (0.0, 0.001, 0.005, 0.01, 0.02)
RGB_THRESHOLDS = (0.0, 0.01, 0.05, 0.1, 0.2, 0.5, 1.0)
def _coords(h, w, device):
y, x = torch.meshgrid(
torch.arange(h, device=device), torch.arange(w, device=device), indexing="ij"
)
return torch.stack((x, y), -1).reshape(-1, 2).float()
@torch.no_grad()
def visible_features(
source: Tensor,
target: Tensor,
adapted: Tensor,
rgb: Tensor,
readout: Readout,
state: FastState,
z: Tensor,
) -> dict[str, float]:
"""Extract observed-only signals; no geometry, masks, or task labels enter."""
for name, value in (("source", source), ("target", target), ("adapted", adapted)):
if value.ndim != 3 or min(value.shape) <= 0 or not torch.isfinite(value).all():
raise ValueError(f"{name} must be a finite HWC tensor")
if target.shape != adapted.shape or source.shape[-1] != target.shape[-1]:
raise ValueError("Descriptor shapes do not align")
h, w, _ = target.shape
if rgb.shape != (h, w, 3) or z.shape[:2] != (h, w):
raise ValueError("RGB and projected features must align with target")
before = readout.predict(z) - rgb
after = readout.predict(materialize(z, state)) - rgb
yy, xx = torch.meshgrid(
torch.arange(h, device=target.device), torch.arange(w, device=target.device), indexing="ij"
)
probe = (yy + 2 * xx) % 4 == 3
if not probe.any():
raise ValueError("RGB selection partition 3 is empty")
result = {}
for prefix, a, b in (("rgb", before, after), ("probe", before[probe], after[probe])):
old, new = float(a.square().mean()), float(b.square().mean())
result.update(
{
f"{prefix}_mse_before": old,
f"{prefix}_mse_after": new,
f"{prefix}_decrease": old - new,
f"{prefix}_relative_decrease": (old - new) / max(old, 1e-12),
}
)
for channel in range(3):
result[f"{prefix}_residual_abs_before_{channel}"] = float(a[..., channel].abs().mean())
result[f"{prefix}_residual_abs_after_{channel}"] = float(b[..., channel].abs().mean())
result[f"{prefix}_residual_variance_before_{channel}"] = float(
a[..., channel].var(unbiased=False)
)
result[f"{prefix}_residual_variance_after_{channel}"] = float(
b[..., channel].var(unbiased=False)
)
src = F.normalize(source.reshape(-1, source.shape[-1]), dim=-1)
dst = F.normalize(adapted.reshape(-1, adapted.shape[-1]), dim=-1)
# Target -> frozen source follows the propagation direction without labels.
scores = dst @ src.T
best, index = scores.max(-1)
top = scores.topk(min(2, scores.shape[-1]), dim=-1).values
gap = top[:, 0] - top[:, -1]
probabilities = scores.softmax(-1)
entropy = -(probabilities * probabilities.clamp_min(1e-12).log()).sum(-1)
reverse = scores.argmax(0)
cycle = reverse[index]
dst_xy = _coords(h, w, target.device)
src_xy = _coords(*source.shape[:2], target.device)
src_xy = src_xy / src_xy.new_tensor((max(source.shape[1] - 1, 1), max(source.shape[0] - 1, 1)))
dst_normalized = dst_xy / dst_xy.new_tensor((max(w - 1, 1), max(h - 1, 1)))
displacement = src_xy[index] - dst_normalized
field = displacement.reshape(h, w, 2)
differences = []
if h > 1:
differences.append((field[1:] - field[:-1]).norm(dim=-1).reshape(-1))
if w > 1:
differences.append((field[:, 1:] - field[:, :-1]).norm(dim=-1).reshape(-1))
drift = (adapted - target).norm(dim=-1) / target.norm(dim=-1).clamp_min(1e-12)
result.update(
{
"match_cosine_mean": float(best.mean()),
"match_cosine_median": float(best.median()),
"match_margin_mean": float(gap.mean()),
"match_entropy_mean": float(entropy.mean()),
"mutual_nn_fraction": float(
(cycle == torch.arange(len(dst), device=target.device)).float().mean()
),
"cycle_error": float((dst_normalized[cycle] - dst_normalized).norm(dim=-1).mean()),
"displacement_mean": float(displacement.norm(dim=-1).mean()),
"displacement_std": float(displacement.norm(dim=-1).std(unbiased=False)),
"displacement_smoothness": float(torch.cat(differences).mean()) if differences else 0.0,
"state_a_norm": float(state.a.norm()),
"state_b_norm": float(state.b.norm()),
"state_total_norm": float(torch.sqrt(state.a.square().sum() + state.b.square().sum())),
"feature_drift": float(drift.mean()),
"feature_cosine": float(F.cosine_similarity(target, adapted, dim=-1).mean()),
"feature_large_change_fraction": float((drift > 0.1).float().mean()),
}
)
if not all(np.isfinite(v) for v in result.values()):
raise ValueError("Visible features contain nonfinite values")
return result
def enrich_candidates(features: list[dict], descriptors: list[Tensor] | None = None) -> list[dict]:
"""Add candidate-set statistics; optional descriptors provide exact disagreement."""
if not features:
raise ValueError("Candidates cannot be empty")
values = np.array([f["probe_mse_after"] for f in features], dtype=float)
if not np.isfinite(values).all():
raise ValueError("Candidate RGB losses must be finite")
if descriptors is not None:
if len(descriptors) != len(features) or any(
d.shape != descriptors[0].shape for d in descriptors
):
raise ValueError("Descriptors must align with candidates")
arrays = [F.normalize(d.detach().float(), dim=-1) for d in descriptors]
disagreement = [
float(torch.stack([(a - b).square().sum(-1).mean() for b in arrays]).mean())
for a in arrays
]
else:
drift = np.array([f.get("feature_drift", 0.0) for f in features])
disagreement = np.abs(drift[:, None] - drift[None, :]).mean(1).tolist()
result = []
for i, item in enumerate(features):
competitors = np.delete(values, i)
result.append(
{
**item,
"rgb_rank": float((values < values[i]).sum()),
"rgb_confidence_margin": float(competitors.min() - values[i])
if len(competitors)
else 0.0,
"candidate_disagreement": disagreement[i],
}
)
return result
def _trajectory_key(trajectory):
if len(trajectory) != 2 or any(a not in ACTIONS for a in trajectory):
raise ValueError("Trajectories must contain two known actions")
return tuple(ACTIONS.index(a) for a in trajectory)
def _off(trajectories):
matches = [i for i, t in enumerate(trajectories) if list(t) == ["OFF", "OFF"]]
if len(matches) != 1 or len({tuple(t) for t in trajectories}) != len(trajectories):
raise ValueError("Candidates require unique trajectories and exactly one OFF/OFF")
for trajectory in trajectories:
_trajectory_key(trajectory)
return matches[0]
def _best(values, trajectories, maximize=True):
return min(
range(len(values)),
key=lambda i: ((-1 if maximize else 1) * values[i], _trajectory_key(trajectories[i])),
)
def _groups(rows):
result = {}
for row in rows:
key = (str(row["sequence"]), str(row["target_slot"]))
result.setdefault(key, []).append(row)
return list(result.values())
def _matrix(rows, names):
x = np.asarray([[r["visible"][n] for n in names] for r in rows], dtype=float)
if x.ndim != 2 or not np.isfinite(x).all():
raise ValueError("Features must be finite and match the locked schema")
return x
def _fit(rows, names, alpha):
x = _matrix(rows, names)
baselines = {}
for group in _groups(rows):
off = _off([r["trajectory"] for r in group])
key = (str(group[0]["sequence"]), str(group[0]["target_slot"]))
baselines[key] = float(group[off]["J"])
y = np.array(
[float(r["J"]) - baselines[(str(r["sequence"]), str(r["target_slot"]))] for r in rows],
dtype=float,
)
mean, scale = x.mean(0), x.std(0)
scale = np.where(scale > 1e-12, scale, 1.0)
normalized = (x - mean) / scale
intercept = float(y.mean())
coefficients = np.linalg.solve(
normalized.T @ normalized + alpha * np.eye(len(names)), normalized.T @ (y - intercept)
)
return {
"feature_names": names,
"mean": mean.tolist(),
"scale": scale.tolist(),
"coefficients": coefficients.tolist(),
"intercept": intercept,
"alpha": float(alpha),
}
def _pred(rows, model):
return (_matrix(rows, model["feature_names"]) - np.array(model["mean"])) / np.array(
model["scale"]
) @ np.array(model["coefficients"]) + model["intercept"]
def _choose_visible(predictions, trajectories, threshold):
off = _off(trajectories)
predictions = np.asarray(predictions, dtype=float).copy()
predictions[off] = 0.0
best = _best(predictions, trajectories)
return best if predictions[best] > threshold else off
def _choose_rgb(features, trajectories, threshold):
off = _off(trajectories)
losses = [f["probe_mse_after"] for f in features]
best = _best(losses, trajectories, maximize=False)
relative = (losses[off] - losses[best]) / max(losses[off], 1e-12)
return best if relative > threshold else off
def _reward(group, choice):
off = _off([r["trajectory"] for r in group])
return float(group[choice]["J"] - group[off]["J"])
def _sequence_mean(records):
sequences = {}
for sequence, reward in records:
sequences.setdefault(sequence, []).append(reward)
return float(np.mean([np.mean(v) for v in sequences.values()]))
def fit_selectors(rows: list[dict], config: dict) -> dict:
"""Fit only discovery rows; reject rather than silently consume other splits.
Ridge hyperparameters use leave-one-sequence-out model predictions. RGB
thresholds are searched directly on discovery rewards; their tuning score
is not an out-of-fold generalization estimate.
"""
if not rows or any(r.get("split") != "discovery" for r in rows):
raise ValueError("Selector fitting accepts discovery rows only")
if any(not np.isfinite(float(r["J"])) for r in rows):
raise ValueError("Discovery rewards must be finite")
output = {
"schema_version": SCHEMA,
"fit_split": "discovery",
"training_sequences": sorted({str(r["sequence"]) for r in rows}),
"by_resolution": {},
"metadata": dict(config),
}
for resolution in sorted({str(r["resolution"]) for r in rows}):
selected = [r for r in rows if str(r["resolution"]) == resolution]
sequences = sorted({str(r["sequence"]) for r in selected})
if len(sequences) < 2:
raise ValueError("LOSO requires at least two discovery sequences per resolution")
names = sorted(selected[0]["visible"])
if not names or any(set(r["visible"]) != set(names) for r in selected):
raise ValueError("Visible feature schema must be consistent")
groups = _groups(selected)
for group in groups:
_off([r["trajectory"] for r in group])
choices = []
for alpha in ALPHAS:
heldout = []
for sequence in sequences:
model = _fit([r for r in selected if str(r["sequence"]) != sequence], names, alpha)
heldout.extend(
(g, _pred(g, model)) for g in groups if str(g[0]["sequence"]) == sequence
)
for threshold in THRESHOLDS:
rewards = [
(
str(g[0]["sequence"]),
_reward(g, _choose_visible(p, [r["trajectory"] for r in g], threshold)),
)
for g, p in heldout
]
choices.append((_sequence_mean(rewards), alpha, threshold))
# Prefer more abstention and stronger regularization when rewards tie.
cv_reward, alpha, threshold = max(choices, key=lambda x: (x[0], x[2], x[1]))
rgb_choices = []
for t in RGB_THRESHOLDS:
records = [
(
str(g[0]["sequence"]),
_reward(
g, _choose_rgb([r["visible"] for r in g], [r["trajectory"] for r in g], t)
),
)
for g in groups
]
rgb_choices.append((_sequence_mean(records), t))
rgb_reward, rgb_threshold = max(rgb_choices)
trajectories = sorted({tuple(r["trajectory"]) for r in selected}, key=_trajectory_key)
if any({tuple(r["trajectory"]) for r in g} != set(trajectories) for g in groups):
raise ValueError("Every discovery target must contain the same candidate set")
fixed_rewards = []
for t in trajectories:
records = [
(
str(g[0]["sequence"]),
_reward(g, next(i for i, r in enumerate(g) if tuple(r["trajectory"]) == t)),
)
for g in groups
]
fixed_rewards.append(_sequence_mean(records))
best_fixed = trajectories[_best(fixed_rewards, trajectories)]
artifact = {
**_fit(selected, names, alpha),
"schema_version": SCHEMA,
"fit_split": "discovery",
"resolution": resolution,
"training_sequences": sequences,
"threshold": threshold,
"rgb_threshold": rgb_threshold,
"discovery_best_fixed": list(best_fixed),
"metadata": {
"cv": "leave_one_sequence_out",
"reward": "sequence_mean_J_candidate_minus_J_OFF",
"regression_target": "J_candidate_minus_same_pair_J_OFF",
"off_prediction": "forced_zero_before_ranking",
"cv_reward": cv_reward,
"rgb_discovery_tuning_reward": rgb_reward,
"rgb_threshold_selection": "discovery_reward_search_not_oof",
"alpha_grid": list(ALPHAS),
"threshold_grid": list(THRESHOLDS),
"rgb_threshold_grid": list(RGB_THRESHOLDS),
"config": dict(config),
},
}
output["by_resolution"][resolution] = artifact
return output
def validate_artifact(artifact: dict) -> dict:
"""Gate deployment on a complete discovery-fitted artifact."""
if artifact.get("schema_version") != SCHEMA or artifact.get("fit_split") != "discovery":
raise ValueError("Expected a locked discovery selector artifact")
if "by_resolution" in artifact:
if len(artifact["by_resolution"]) != 1:
raise ValueError("Select the resolution-specific artifact before inference")
artifact = next(iter(artifact["by_resolution"].values()))
return validate_artifact(artifact)
names = artifact.get("feature_names", [])
if not names or len(set(names)) != len(names) or not artifact.get("training_sequences"):
raise ValueError("Missing feature schema or training provenance")
for name in ("coefficients", "mean", "scale"):
array = np.asarray(artifact.get(name), dtype=float)
if array.shape != (len(names),) or not np.isfinite(array).all():
raise ValueError(f"Invalid artifact {name}")
if np.any(np.asarray(artifact["scale"]) <= 0):
raise ValueError("Artifact scales must be positive")
for name in ("intercept", "alpha", "threshold", "rgb_threshold"):
if name not in artifact or not np.isfinite(artifact[name]):
raise ValueError(f"Invalid artifact {name}")
if artifact["threshold"] < 0 or artifact["rgb_threshold"] < 0 or artifact["alpha"] <= 0:
raise ValueError("Invalid selector hyperparameters")
_trajectory_key(artifact["discovery_best_fixed"])
return artifact
def predict_candidates(features: list[dict], artifact: dict) -> list[float]:
"""Return model-predicted J-minus-OFF reward without labels.
These are raw regression outputs; selection pins the known OFF reward to
zero before ranking and applying the abstention threshold.
"""
model = validate_artifact(artifact)
if not features or any(set(f) != set(model["feature_names"]) for f in features):
raise ValueError("Inference features must match the locked schema exactly")
return _pred([{"visible": f} for f in features], model).tolist()
def select_candidates(features: list[dict], trajectories: list[list[str]], artifact: dict) -> dict:
model = validate_artifact(artifact)
if len(features) != len(trajectories):
raise ValueError("Features and trajectories must align")
predictions = predict_candidates(features, model)
_off(trajectories)
fixed = [i for i, t in enumerate(trajectories) if list(t) == model["discovery_best_fixed"]]
if len(fixed) != 1:
raise ValueError("Locked best-fixed trajectory is absent")
return {
"rgb_selector": _choose_rgb(features, trajectories, model["rgb_threshold"]),
"visible_selector": _choose_visible(predictions, trajectories, model["threshold"]),
"best_fixed": fixed[0],
}