src/mcss/measurement_planning/information.py
"""Information contract of MEASUREMENT_PLANNING_TASK_V1: what a deployable selector may see.
An episode is one record and one context role (A or B). A deployable selector receives a
MeasurementEnv built from a sanitized copy of the record that holds only that role's context frames
and the candidate frames, so query frames (RGB, depth, camera) and the other role's context frames
cannot be reached through it. Candidate depth is returned only by request_depth, within the budget,
and every read is logged with its purpose. The teacher (oracle) path is a separate object that is
never handed to a selector.
"""
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass
import numpy as np
import torch
from mcss.dynamic.types import OnlineObservation
from mcss.mechanism_pilot.direct_capacity_contracts import _Media
from mcss.mechanism_pilot.rgbd_evidence_carrier import RGBDContext
from mcss.types import Cameras
ROLE_KEYS = {"A": "context_a", "B": "context_b"}
QUERY_KEYS = ("query", "primary_query", "secondary_query")
ROLE_FIELDS = (
"context_a",
"context_b",
"query",
"primary_query",
"secondary_query",
"stream",
"far_query",
"ignored",
)
CONTEXT_PURPOSE = "CONTEXT_ONLY_RGBD"
CANDIDATE_RGB_PURPOSE = "CANDIDATE_RGB_CAMERA"
CANDIDATE_DEPTH_PURPOSE = "CANDIDATE_DEPTH_REQUESTED"
TEACHER_PURPOSE = "ORACLE_TEACHER_ONLY"
QUERY_PURPOSE = "SEALED_QUERY_EVALUATION"
QUERY_CAMERA_PURPOSE = "QUERY_CAMERA_TEACHER_ONLY"
def query_ids(record):
"""Every frame any query role of the record names (primary, secondary and plain query lists)."""
roles = record["roles"]
return sorted(set().union(*(set(roles.get(k, [])) for k in QUERY_KEYS)))
def candidate_pool(record, min_query_gap=0):
"""Frames no role uses, in frame order; with min_query_gap > 0, frames whose index is closer
than that to any query frame are left out. A candidate is never a context frame of either role
nor a query."""
roles = record["roles"]
used = set().union(*(set(roles.get(k, [])) for k in ROLE_FIELDS))
queries = query_ids(record)
return [
f["frame_id"]
for f in record["frames"]
if f["frame_id"] not in used
and all(abs(f["frame_id"] - q) >= min_query_gap for q in queries)
]
def sanitized_record(record, role, pool):
"""The only record a selector's loader holds: the role's context and the candidate frames."""
context = list(record["roles"][ROLE_KEYS[role]])
forbidden = set(query_ids(record)) | set(record["roles"].get("far_query", []))
if set(context) & set(pool) or set(pool) & forbidden or set(context) & forbidden:
raise PermissionError("Candidates must be disjoint from context and query frames")
keep = set(context) | set(pool)
return {
"scene_id": record["scene_id"],
"frames": [deepcopy(f) for f in record["frames"] if f["frame_id"] in keep],
"roles": {"context": context, "candidates": list(pool)},
}
def single_camera(batch, j):
return Cameras(
batch.cameras.intrinsics[0, j], batch.cameras.c2w[0, j], batch.cameras.image_size
)
def to_context(batch):
"""ObservationBatch with depth -> RGBDContext, exactly as the frozen context_of builds it."""
observations = [
OnlineObservation(batch.scene_id, fid, batch.rgb[0, j], single_camera(batch, j))
for j, fid in enumerate(batch.frame_ids)
]
return RGBDContext(observations, batch.depth[0, :, 0])
@dataclass(frozen=True)
class CandidateView:
"""What arrives with a candidate before any request: RGB [3, H, W] in [0, 1] and its camera."""
frame_id: int
rgb: torch.Tensor
camera: Cameras
class BudgetExceededError(PermissionError):
pass
class MeasurementEnv:
"""The deployable interface: measured context, candidate RGB and cameras, requested depth
only."""
def __init__(self, record, role, pool, budget, image_size, base, access_log):
if not 0 <= budget <= len(pool):
raise ValueError("Budget must lie within the candidate pool")
self.scene_id, self.role = record["scene_id"], role
self._record = sanitized_record(record, role, pool)
self._media = _Media(self._record, image_size, base, "cpu", access_log)
self.context_ids = tuple(self._record["roles"]["context"])
self.candidate_ids = tuple(pool)
self.budget = int(budget)
self._requested = []
self._context = None
self._candidates = None
def context(self):
"""The role's measured RGB-D context (fused before any decision)."""
if self._context is None:
batch = self._media.read(self.context_ids, depth_allowed=True, purpose=CONTEXT_PURPOSE)
self._context = to_context(batch)
return self._context
def candidates(self):
"""Candidate RGB and cameras; no depth."""
if self._candidates is None:
batch = self._media.read(
self.candidate_ids, depth_allowed=False, purpose=CANDIDATE_RGB_PURPOSE
)
self._candidates = tuple(
CandidateView(fid, batch.rgb[0, j], single_camera(batch, j))
for j, fid in enumerate(batch.frame_ids)
)
return self._candidates
def request_depth(self, frame_id):
"""REQUEST_DEPTH: spends one unit of budget and returns that candidate as a one-frame
RGBDContext."""
if frame_id not in self.candidate_ids:
raise PermissionError("Depth may be requested only for a candidate frame")
if frame_id in self._requested:
raise ValueError("Each candidate's depth is requested at most once")
if len(self._requested) >= self.budget:
raise BudgetExceededError("Depth budget exhausted")
self._requested.append(frame_id)
batch = self._media.read((frame_id,), depth_allowed=True, purpose=CANDIDATE_DEPTH_PURPOSE)
return to_context(batch)
@property
def requested(self):
return tuple(self._requested)
@property
def remaining(self):
return self.budget - len(self._requested)
class TeacherAccess:
"""ORACLE / TEACHER ONLY: full access to one record for oracle utilities and scoring. Never a
selector input."""
def __init__(self, record, query_frames, image_size, base, access_log):
self._record = deepcopy(record)
self._media = _Media(self._record, image_size, base, "cpu", access_log)
self.query_frames = tuple(query_frames)
self.image_size = tuple(image_size)
def frames(self, ids):
"""Measured RGB-D of context or candidate frames, for the oracle's exhaustive subset
evaluation."""
if set(ids) & set(self.query_frames) or set(ids) & set(query_ids(self._record)):
raise PermissionError("Query frames are scored, never fused")
return to_context(self._media.read(tuple(ids), depth_allowed=True, purpose=TEACHER_PURPOSE))
def query_camera(self, frame_id):
"""A query camera for scoring (and the query-aware diagnostic only), as the frozen
query_camera builds it."""
if frame_id not in self.query_frames:
raise PermissionError("Not a query frame of this episode")
frame = self._media.frames[frame_id]
self._media.access_log.append(
{
"scene_id": self._record["scene_id"],
"frame_id": frame_id,
"purpose": QUERY_CAMERA_PURPOSE,
"channels": ["camera"],
}
)
return Cameras(
torch.tensor(frame["intrinsics"], dtype=torch.float32),
torch.tensor(frame["c2w"], dtype=torch.float32),
self.image_size,
)
def query_depth(self, frame_id):
"""Query ground truth as the arms evaluator scores it (np.load, squeeze, float64)."""
if frame_id not in self.query_frames:
raise PermissionError("Not a query frame of this episode")
frame = self._media.frames[frame_id]
self._media.access_log.append(
{
"scene_id": self._record["scene_id"],
"frame_id": frame_id,
"purpose": QUERY_PURPOSE,
"channels": ["depth"],
}
)
return np.load(self._media._path(frame, "depth")).squeeze().astype(np.float64)